


<!DOCTYPE html>
<!--[if IE 8]><html class="no-js lt-ie9" lang="en" > <![endif]-->
<!--[if gt IE 8]><!--> <html class="no-js" lang="en" > <!--<![endif]-->
<head>
  <meta charset="utf-8">
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  
  <title>torch.distributed.distributed_c10d &mdash; PyTorch master documentation</title>
  

  
  
  
  
    <link rel="canonical" href="https://pytorch.org/docs/stable/_modules/torch/distributed/distributed_c10d.html"/>
  

  

  
  
    

  

  <link rel="stylesheet" href="../../../_static/css/theme.css" type="text/css" />
  <!-- <link rel="stylesheet" href="../../../_static/pygments.css" type="text/css" /> -->
  <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.10.0-beta/dist/katex.min.css" type="text/css" />
  <link rel="stylesheet" href="../../../_static/css/jit.css" type="text/css" />
  <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.11.1/dist/katex.min.css" type="text/css" />
  <link rel="stylesheet" href="../../../_static/katex-math.css" type="text/css" />
    <link rel="index" title="Index" href="../../../genindex.html" />
    <link rel="search" title="Search" href="../../../search.html" /> 

  
  <script src="../../../_static/js/modernizr.min.js"></script>

  <!-- Preload the theme fonts -->

<link rel="preload" href="../../../_static/fonts/FreightSans/freight-sans-book.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="../../../_static/fonts/FreightSans/freight-sans-medium.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="../../../_static/fonts/IBMPlexMono/IBMPlexMono-Medium.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="../../../_static/fonts/FreightSans/freight-sans-bold.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="../../../_static/fonts/FreightSans/freight-sans-medium-italic.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="../../../_static/fonts/IBMPlexMono/IBMPlexMono-SemiBold.woff2" as="font" type="font/woff2" crossorigin="anonymous">

<!-- Preload the katex fonts -->

<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Math-Italic.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Main-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Main-Bold.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size1-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size4-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size2-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size3-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Caligraphic-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
</head>

<div class="container-fluid header-holder tutorials-header" id="header-holder">
  <div class="container">
    <div class="header-container">
      <a class="header-logo" href="https://pytorch.org/" aria-label="PyTorch"></a>

      <div class="main-menu">
        <ul>
          <li>
            <a href="https://pytorch.org/get-started">Get Started</a>
          </li>

          <li>
            <div class="ecosystem-dropdown">
              <a id="dropdownMenuButton" data-toggle="ecosystem-dropdown">
                Ecosystem
              </a>
              <div class="ecosystem-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/hub"">
                  <span class=dropdown-title>Models (Beta)</span>
                  <p>Discover, publish, and reuse pre-trained models</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/ecosystem">
                  <span class=dropdown-title>Tools & Libraries</span>
                  <p>Explore the ecosystem of tools and libraries</p>
                </a>
              </div>
            </div>
          </li>

          <li>
            <a href="https://pytorch.org/mobile">Mobile</a>
          </li>

          <li>
            <a href="https://pytorch.org/blog/">Blog</a>
          </li>

          <li>
            <a href="https://pytorch.org/tutorials">Tutorials</a>
          </li>

          <li class="active">
            <a href="https://pytorch.org/docs/stable/index.html">Docs</a>
          </li>

          <li>
            <div class="resources-dropdown">
              <a id="resourcesDropdownButton" data-toggle="resources-dropdown">
                Resources
              </a>
              <div class="resources-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/resources"">
                  <span class=dropdown-title>Developer Resources</span>
                  <p>Find resources and get questions answered</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/features">
                  <span class=dropdown-title>About</span>
                  <p>Learn about PyTorch’s features and capabilities</p>
                </a>
              </div>
            </div>
          </li>

          <li>
            <a href="https://github.com/pytorch/pytorch">Github</a>
          </li>
        </ul>
      </div>

      <a class="main-menu-open-button" href="#" data-behavior="open-mobile-menu"></a>
    </div>

  </div>
</div>


<body class="pytorch-body">

   

    

    <div class="table-of-contents-link-wrapper">
      <span>Table of Contents</span>
      <a href="#" class="toggle-table-of-contents" data-behavior="toggle-table-of-contents"></a>
    </div>

    <nav data-toggle="wy-nav-shift" class="pytorch-left-menu" id="pytorch-left-menu">
      <div class="pytorch-side-scroll">
        <div class="pytorch-menu pytorch-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          <div class="pytorch-left-menu-search">
            

            
              
              
                <div class="version">
                  master (1.5.0 )
                </div>
              
            

            


  


<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search Docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>

            
          </div>

          
<div>
  <a style="color:#F05732" href="https://pytorch.org/docs/stable/_modules/torch/distributed/distributed_c10d.html">
    You are viewing unstable developer preview docs.
    Click here to view docs for latest stable release.
  </a>
</div>

            
            
              
            
            
              <p class="caption"><span class="caption-text">Notes</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/amp_examples.html">Automatic Mixed Precision examples</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/autograd.html">Autograd mechanics</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/broadcasting.html">Broadcasting semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/cpu_threading_torchscript_inference.html">CPU threading and TorchScript inference</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/cuda.html">CUDA semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/ddp.html">Distributed Data Parallel</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/extending.html">Extending PyTorch</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/faq.html">Frequently Asked Questions</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/large_scale_deployments.html">Features for large-scale deployments</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/multiprocessing.html">Multiprocessing best practices</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/randomness.html">Reproducibility</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/serialization.html">Serialization semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../notes/windows.html">Windows FAQ</a></li>
</ul>
<p class="caption"><span class="caption-text">Language Bindings</span></p>
<ul>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/cppdocs/">C++ API</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../packages.html">Javadoc</a></li>
</ul>
<p class="caption"><span class="caption-text">Python API</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../torch.html">torch</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../nn.html">torch.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../nn.functional.html">torch.nn.functional</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../tensors.html">torch.Tensor</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../tensor_attributes.html">Tensor Attributes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../tensor_view.html">Tensor Views</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../autograd.html">torch.autograd</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../cuda.html">torch.cuda</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../amp.html">torch.cuda.amp</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../distributed.html">torch.distributed</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../distributions.html">torch.distributions</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../hub.html">torch.hub</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../jit.html">torch.jit</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../nn.init.html">torch.nn.init</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../onnx.html">torch.onnx</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../optim.html">torch.optim</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../quantization.html">Quantization</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../rpc/index.html">Distributed RPC Framework</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../random.html">torch.random</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../sparse.html">torch.sparse</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../storage.html">torch.Storage</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../bottleneck.html">torch.utils.bottleneck</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../checkpoint.html">torch.utils.checkpoint</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../cpp_extension.html">torch.utils.cpp_extension</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../data.html">torch.utils.data</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../dlpack.html">torch.utils.dlpack</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../model_zoo.html">torch.utils.model_zoo</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../tensorboard.html">torch.utils.tensorboard</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../type_info.html">Type Info</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../named_tensor.html">Named Tensors</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../name_inference.html">Named Tensors operator coverage</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../__config__.html">torch.__config__</a></li>
</ul>
<p class="caption"><span class="caption-text">Libraries</span></p>
<ul>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/audio">torchaudio</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/text">torchtext</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/elastic/">TorchElastic</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/serve">TorchServe</a></li>
<li class="toctree-l1"><a class="reference external" href="http://pytorch.org/xla/">PyTorch on XLA Devices</a></li>
</ul>
<p class="caption"><span class="caption-text">Community</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../community/contribution_guide.html">PyTorch Contribution Guide</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../community/governance.html">PyTorch Governance</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../community/persons_of_interest.html">PyTorch Governance | Persons of Interest</a></li>
</ul>

            
          

        </div>
      </div>
    </nav>

    <div class="pytorch-container">
      <div class="pytorch-page-level-bar" id="pytorch-page-level-bar">
        <div class="pytorch-breadcrumbs-wrapper">
          















<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="pytorch-breadcrumbs">
    
      <li>
        <a href="../../../index.html">
          
            Docs
          
        </a> &gt;
      </li>

        
          <li><a href="../../index.html">Module code</a> &gt;</li>
        
          <li><a href="../../torch.html">torch</a> &gt;</li>
        
          <li><a href="../distributed.html">torch.distributed</a> &gt;</li>
        
      <li>torch.distributed.distributed_c10d</li>
    
    
      <li class="pytorch-breadcrumbs-aside">
        
      </li>
    
  </ul>

  
</div>
        </div>

        <div class="pytorch-shortcuts-wrapper" id="pytorch-shortcuts-wrapper">
          Shortcuts
        </div>
      </div>

      <section data-toggle="wy-nav-shift" id="pytorch-content-wrap" class="pytorch-content-wrap">
        <div class="pytorch-content-left">

        
          
          <div class="rst-content">
          
            <div role="main" class="main-content" itemscope="itemscope" itemtype="http://schema.org/Article">
             <article itemprop="articleBody" id="pytorch-article" class="pytorch-article">
              
  <h1>Source code for torch.distributed.distributed_c10d</h1><div class="highlight"><pre>
<span></span><span class="kn">import</span> <span class="nn">torch</span>
<span class="kn">import</span> <span class="nn">warnings</span>
<span class="kn">from</span> <span class="nn">torch._six</span> <span class="kn">import</span> <span class="n">string_classes</span>
<span class="kn">from</span> <span class="nn">datetime</span> <span class="kn">import</span> <span class="n">timedelta</span>

<span class="c1"># This module is wildcard imported from torch.distributed.</span>
<span class="c1"># TODO: specify __all__</span>

<span class="kn">from</span> <span class="nn">.constants</span> <span class="kn">import</span> <span class="n">default_pg_timeout</span>
<span class="kn">from</span> <span class="nn">.rendezvous</span> <span class="kn">import</span> <span class="n">rendezvous</span><span class="p">,</span> <span class="n">register_rendezvous_handler</span>  <span class="c1"># noqa: F401</span>
<span class="kn">from</span> <span class="nn">.</span> <span class="kn">import</span> <span class="p">(</span>
    <span class="n">AllreduceOptions</span><span class="p">,</span>
    <span class="n">AllreduceCoalescedOptions</span><span class="p">,</span>
    <span class="n">BroadcastOptions</span><span class="p">,</span>
    <span class="n">GatherOptions</span><span class="p">,</span>
    <span class="n">ReduceOptions</span><span class="p">,</span>
    <span class="n">ReduceScatterOptions</span><span class="p">,</span>
    <span class="n">ScatterOptions</span><span class="p">,</span>
<span class="p">)</span>
<span class="kn">from</span> <span class="nn">.</span> <span class="kn">import</span> <span class="n">ReduceOp</span>
<span class="kn">from</span> <span class="nn">.</span> <span class="kn">import</span> <span class="n">PrefixStore</span>


<span class="n">_MPI_AVAILABLE</span> <span class="o">=</span> <span class="kc">True</span>
<span class="n">_NCCL_AVAILABLE</span> <span class="o">=</span> <span class="kc">True</span>
<span class="n">_GLOO_AVAILABLE</span> <span class="o">=</span> <span class="kc">True</span>


<span class="k">try</span><span class="p">:</span>
    <span class="n">from</span><span class="o">.</span> <span class="kn">import</span> <span class="nn">ProcessGroupMPI</span>
<span class="k">except</span> <span class="ne">ImportError</span><span class="p">:</span>
    <span class="n">_MPI_AVAILABLE</span> <span class="o">=</span> <span class="kc">False</span>

<span class="k">try</span><span class="p">:</span>
    <span class="n">from</span><span class="o">.</span> <span class="kn">import</span> <span class="nn">ProcessGroupNCCL</span>
<span class="k">except</span> <span class="ne">ImportError</span><span class="p">:</span>
    <span class="n">_NCCL_AVAILABLE</span> <span class="o">=</span> <span class="kc">False</span>

<span class="k">try</span><span class="p">:</span>
    <span class="n">from</span><span class="o">.</span> <span class="kn">import</span> <span class="nn">ProcessGroupGloo</span>
<span class="k">except</span> <span class="ne">ImportError</span><span class="p">:</span>
    <span class="n">_GLOO_AVAILABLE</span> <span class="o">=</span> <span class="kc">False</span>


<div class="viewcode-block" id="Backend"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.Backend">[docs]</a><span class="k">class</span> <span class="nc">Backend</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    An enum-like class of available backends: GLOO, NCCL, and MPI.</span>

<span class="sd">    The values of this class are lowercase strings, e.g., ``&quot;gloo&quot;``. They can</span>
<span class="sd">    be accessed as attributes, e.g., ``Backend.NCCL``.</span>

<span class="sd">    This class can be directly called to parse the string, e.g.,</span>
<span class="sd">    ``Backend(backend_str)`` will check if ``backend_str`` is valid, and</span>
<span class="sd">    return the parsed lowercase string if so. It also accepts uppercase strings,</span>
<span class="sd">    e.g., ``Backend(&quot;GLOO&quot;)`` returns ``&quot;gloo&quot;``.</span>

<span class="sd">    .. note:: The entry ``Backend.UNDEFINED`` is present but only used as</span>
<span class="sd">              initial value of some fields. Users should neither use it directly</span>
<span class="sd">              nor assume its existence.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">UNDEFINED</span> <span class="o">=</span> <span class="s2">&quot;undefined&quot;</span>
    <span class="n">GLOO</span> <span class="o">=</span> <span class="s2">&quot;gloo&quot;</span>
    <span class="n">NCCL</span> <span class="o">=</span> <span class="s2">&quot;nccl&quot;</span>
    <span class="n">MPI</span> <span class="o">=</span> <span class="s2">&quot;mpi&quot;</span>
    <span class="n">TCP</span> <span class="o">=</span> <span class="s2">&quot;tcp&quot;</span>

    <span class="k">def</span> <span class="fm">__new__</span><span class="p">(</span><span class="bp">cls</span><span class="p">,</span> <span class="n">name</span><span class="p">):</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">string_classes</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Backend name must be a string, but got: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">name</span><span class="p">))</span>
        <span class="n">value</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">Backend</span><span class="p">,</span> <span class="n">name</span><span class="o">.</span><span class="n">upper</span><span class="p">(),</span> <span class="n">Backend</span><span class="o">.</span><span class="n">UNDEFINED</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">value</span> <span class="o">==</span> <span class="n">Backend</span><span class="o">.</span><span class="n">TCP</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;TCP backend has been deprecated. Please use &quot;</span>
                             <span class="s2">&quot;Gloo or MPI backend for collective operations &quot;</span>
                             <span class="s2">&quot;on CPU tensors.&quot;</span><span class="p">)</span>
        <span class="k">elif</span> <span class="n">value</span> <span class="o">==</span> <span class="n">Backend</span><span class="o">.</span><span class="n">UNDEFINED</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Invalid backend: &#39;</span><span class="si">{}</span><span class="s2">&#39;&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">name</span><span class="p">))</span>
        <span class="k">return</span> <span class="n">value</span></div>

<span class="c1"># `_backend`, `dist_backend`, and `reduce_op` are here to maintain backward</span>
<span class="c1"># compatibility with pre-c10d distributed package.</span>
<span class="c1"># TODO: remove them when users are ready to take a hard dependency on PyTorch 1.</span>
<span class="n">_backend</span> <span class="o">=</span> <span class="n">Backend</span><span class="o">.</span><span class="n">UNDEFINED</span>
<span class="n">dist_backend</span> <span class="o">=</span> <span class="n">Backend</span>


<div class="viewcode-block" id="reduce_op"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.reduce_op">[docs]</a><span class="k">class</span> <span class="nc">reduce_op</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Deprecated enum-like class for reduction operations: ``SUM``, ``PRODUCT``,</span>
<span class="sd">    ``MIN``, and ``MAX``.</span>

<span class="sd">    :class:`~torch.distributed.ReduceOp` is recommended to use instead.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="c1"># __members__ is a dict storing key-value pairs for enum classes</span>
        <span class="k">for</span> <span class="n">k</span><span class="p">,</span> <span class="n">v</span> <span class="ow">in</span> <span class="n">ReduceOp</span><span class="o">.</span><span class="n">__members__</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
            <span class="nb">setattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">k</span><span class="p">,</span> <span class="n">v</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">__members__</span> <span class="o">=</span> <span class="n">ReduceOp</span><span class="o">.</span><span class="n">__members__</span>

    <span class="k">def</span> <span class="fm">__getattribute__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">key</span><span class="p">):</span>
        <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="s2">&quot;torch.distributed.reduce_op is deprecated, please use &quot;</span>
                      <span class="s2">&quot;torch.distributed.ReduceOp instead&quot;</span><span class="p">)</span>
        <span class="k">return</span> <span class="nb">object</span><span class="o">.</span><span class="fm">__getattribute__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">key</span><span class="p">)</span></div>

<span class="n">reduce_op</span> <span class="o">=</span> <span class="n">reduce_op</span><span class="p">()</span>


<span class="k">class</span> <span class="nc">group</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
    <span class="n">WORLD</span> <span class="o">=</span> <span class="nb">object</span><span class="p">()</span>


<span class="k">class</span> <span class="nc">GroupMember</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
    <span class="c1"># Alias to group.WORLD for backward compatibility</span>
    <span class="n">WORLD</span> <span class="o">=</span> <span class="n">group</span><span class="o">.</span><span class="n">WORLD</span>
    <span class="n">NON_GROUP_MEMBER</span> <span class="o">=</span> <span class="nb">object</span><span class="p">()</span>


<span class="c1"># Cached process groups</span>
<span class="c1"># For NCCL and GLOO pg, it is a map from ProcessGroup to (Backend, Store)</span>
<span class="c1"># For MPI pg, it is a map from ProcessGroup to (Backend, None)</span>
<span class="n">_pg_map</span> <span class="o">=</span> <span class="p">{}</span>
<span class="c1"># Process group&#39;s names, map from ProcessGroup to str</span>
<span class="n">_pg_names</span> <span class="o">=</span> <span class="p">{}</span>
<span class="c1"># Process group&#39;s global rank to local rank mapping</span>
<span class="n">_pg_group_ranks</span> <span class="o">=</span> <span class="p">{}</span>

<span class="c1"># Default process group state</span>
<span class="n">_default_pg</span> <span class="o">=</span> <span class="kc">None</span>
<span class="n">_default_pg_init_method</span> <span class="o">=</span> <span class="kc">None</span>

<span class="c1"># Process group count for default naming</span>
<span class="n">_group_count</span> <span class="o">=</span> <span class="mi">0</span>


<span class="k">def</span> <span class="nf">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Helper that checks if the current process&#39;s rank is not in a given group</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="k">return</span> <span class="kc">False</span>
    <span class="k">return</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">NON_GROUP_MEMBER</span>


<span class="k">def</span> <span class="nf">_get_group_rank</span><span class="p">(</span><span class="n">group</span><span class="p">,</span> <span class="n">rank</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Helper that gets a given group&#39;s local rank in the group from a given global</span>
<span class="sd">    rank</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">group</span> <span class="ow">is</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;group.WORLD does not have local rank to global &quot;</span>
                           <span class="s2">&quot;rank mapping&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">group</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">_pg_group_ranks</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;The given group does not exist&quot;</span><span class="p">)</span>
    <span class="k">try</span><span class="p">:</span>
        <span class="n">group_rank</span> <span class="o">=</span> <span class="n">_pg_group_ranks</span><span class="p">[</span><span class="n">group</span><span class="p">][</span><span class="n">rank</span><span class="p">]</span>
    <span class="k">except</span> <span class="ne">KeyError</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;The global rank is not part of the group&quot;</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">group_rank</span>


<span class="k">def</span> <span class="nf">_get_global_rank</span><span class="p">(</span><span class="n">group</span><span class="p">,</span> <span class="n">group_rank</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Helper that gets a given group&#39;s global rank from a given local rank in the</span>
<span class="sd">    group</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">group</span> <span class="ow">is</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;group.WORLD does not have local rank to global &quot;</span>
                           <span class="s2">&quot;rank mapping&quot;</span><span class="p">)</span>
    <span class="n">group_rank_map</span> <span class="o">=</span> <span class="n">_pg_group_ranks</span><span class="p">[</span><span class="n">group</span><span class="p">]</span>
    <span class="k">for</span> <span class="n">rank</span><span class="p">,</span> <span class="n">grp_rank</span> <span class="ow">in</span> <span class="n">group_rank_map</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
        <span class="k">if</span> <span class="n">grp_rank</span> <span class="o">==</span> <span class="n">group_rank</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">rank</span>
    <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;The group rank is not part of the group&quot;</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_check_default_pg</span><span class="p">():</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Helper that checks if the default ProcessGroup has been initialized, with</span>
<span class="sd">    assertion</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">assert</span> <span class="n">_default_pg</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">,</span> \
        <span class="s2">&quot;Default process group is not initialized&quot;</span>


<span class="k">def</span> <span class="nf">_get_group_size</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Helper that gets a given group&#39;s world size</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">group</span> <span class="ow">is</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_check_default_pg</span><span class="p">()</span>
        <span class="k">return</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">size</span><span class="p">()</span>
    <span class="k">if</span> <span class="n">group</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">_pg_group_ranks</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;The given group does not exist&quot;</span><span class="p">)</span>
    <span class="k">return</span> <span class="nb">len</span><span class="p">(</span><span class="n">_pg_group_ranks</span><span class="p">[</span><span class="n">group</span><span class="p">])</span>


<span class="k">def</span> <span class="nf">_check_single_tensor</span><span class="p">(</span><span class="n">param</span><span class="p">,</span> <span class="n">param_name</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Helper to check that the parameter ``param_name`` is a single tensor.</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">param</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Invalid function argument. Expected parameter `</span><span class="si">{}</span><span class="s2">` &quot;</span>
                           <span class="s2">&quot;to be of type torch.Tensor.&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">param_name</span><span class="p">))</span>


<span class="k">def</span> <span class="nf">_check_tensor_list</span><span class="p">(</span><span class="n">param</span><span class="p">,</span> <span class="n">param_name</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Helper to check that the parameter ``param_name`` is a list of tensors.</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">param</span><span class="p">,</span> <span class="nb">list</span><span class="p">)</span> <span class="ow">or</span> \
       <span class="ow">not</span> <span class="nb">all</span><span class="p">(</span><span class="nb">isinstance</span><span class="p">(</span><span class="n">p</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">)</span> <span class="k">for</span> <span class="n">p</span> <span class="ow">in</span> <span class="n">param</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Invalid function argument. Expected parameter `</span><span class="si">{}</span><span class="s2">` &quot;</span>
                           <span class="s2">&quot;to be of type List[torch.Tensor].&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">param_name</span><span class="p">))</span>


<div class="viewcode-block" id="is_mpi_available"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.is_mpi_available">[docs]</a><span class="k">def</span> <span class="nf">is_mpi_available</span><span class="p">():</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Checks if the MPI backend is available.</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="n">_MPI_AVAILABLE</span></div>


<div class="viewcode-block" id="is_nccl_available"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.is_nccl_available">[docs]</a><span class="k">def</span> <span class="nf">is_nccl_available</span><span class="p">():</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Checks if the NCCL backend is available.</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="n">_NCCL_AVAILABLE</span></div>


<span class="k">def</span> <span class="nf">is_gloo_available</span><span class="p">():</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Checks if the Gloo backend is available.</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="n">_GLOO_AVAILABLE</span>


<div class="viewcode-block" id="is_initialized"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.is_initialized">[docs]</a><span class="k">def</span> <span class="nf">is_initialized</span><span class="p">():</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Checking if the default process group has been initialized</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="n">_default_pg</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span></div>


<span class="k">def</span> <span class="nf">_get_default_group</span><span class="p">():</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Getting the default process group created by init_process_group</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">is_initialized</span><span class="p">():</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Default process group has not been initialized, &quot;</span>
                           <span class="s2">&quot;please make sure to call init_process_group.&quot;</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">_default_pg</span>


<span class="k">def</span> <span class="nf">_get_default_store</span><span class="p">():</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Getting the default store created by init_process_group</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">is_initialized</span><span class="p">():</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Default process group has not been initialized, &quot;</span>
                           <span class="s2">&quot;please make sure to call init_process_group.&quot;</span><span class="p">)</span>
    <span class="n">_</span><span class="p">,</span> <span class="n">default_store</span> <span class="o">=</span> <span class="n">_pg_map</span><span class="p">[</span><span class="n">_default_pg</span><span class="p">]</span>
    <span class="k">return</span> <span class="n">default_store</span>


<div class="viewcode-block" id="get_backend"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.get_backend">[docs]</a><span class="k">def</span> <span class="nf">get_backend</span><span class="p">(</span><span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Returns the backend of the given process group.</span>

<span class="sd">    Arguments:</span>
<span class="sd">        group (ProcessGroup, optional): The process group to work on. The</span>
<span class="sd">            default is the general main process group. If another specific group</span>
<span class="sd">            is specified, the calling process must be part of :attr:`group`.</span>

<span class="sd">    Returns:</span>
<span class="sd">        The backend of the given process group as a lower case string.</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">_check_default_pg</span><span class="p">()</span>

    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">pg</span> <span class="o">=</span> <span class="n">_default_pg</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">pg</span> <span class="o">=</span> <span class="n">group</span>
    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">pg</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Invalid process group specified&quot;</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">_pg_map</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">pg</span><span class="p">,</span> <span class="kc">None</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span></div>


<div class="viewcode-block" id="init_process_group"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.init_process_group">[docs]</a><span class="k">def</span> <span class="nf">init_process_group</span><span class="p">(</span><span class="n">backend</span><span class="p">,</span>
                       <span class="n">init_method</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                       <span class="n">timeout</span><span class="o">=</span><span class="n">default_pg_timeout</span><span class="p">,</span>
                       <span class="n">world_size</span><span class="o">=-</span><span class="mi">1</span><span class="p">,</span>
                       <span class="n">rank</span><span class="o">=-</span><span class="mi">1</span><span class="p">,</span>
                       <span class="n">store</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                       <span class="n">group_name</span><span class="o">=</span><span class="s1">&#39;&#39;</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Initializes the default distributed process group, and this will also</span>
<span class="sd">    initialize the distributed package.</span>

<span class="sd">    There are 2 main ways to initialize a process group:</span>
<span class="sd">        1. Specify ``store``, ``rank``, and ``world_size`` explicitly.</span>
<span class="sd">        2. Specify ``init_method`` (a URL string) which indicates where/how</span>
<span class="sd">           to discover peers. Optionally specify ``rank`` and ``world_size``,</span>
<span class="sd">           or encode all required parameters in the URL and omit them.</span>

<span class="sd">    If neither is specified, ``init_method`` is assumed to be &quot;env://&quot;.</span>


<span class="sd">    Arguments:</span>
<span class="sd">        backend (str or Backend): The backend to use. Depending on</span>
<span class="sd">            build-time configurations, valid values include ``mpi``, ``gloo``,</span>
<span class="sd">            and ``nccl``. This field should be given as a lowercase string</span>
<span class="sd">            (e.g., ``&quot;gloo&quot;``), which can also be accessed via</span>
<span class="sd">            :class:`Backend` attributes (e.g., ``Backend.GLOO``). If using</span>
<span class="sd">            multiple processes per machine with ``nccl`` backend, each process</span>
<span class="sd">            must have exclusive access to every GPU it uses, as sharing GPUs</span>
<span class="sd">            between processes can result in deadlocks.</span>
<span class="sd">        init_method (str, optional): URL specifying how to initialize the</span>
<span class="sd">                                     process group. Default is &quot;env://&quot; if no</span>
<span class="sd">                                     ``init_method`` or ``store`` is specified.</span>
<span class="sd">                                     Mutually exclusive with ``store``.</span>
<span class="sd">        world_size (int, optional): Number of processes participating in</span>
<span class="sd">                                    the job. Required if ``store`` is specified.</span>
<span class="sd">        rank (int, optional): Rank of the current process.</span>
<span class="sd">                              Required if ``store`` is specified.</span>
<span class="sd">        store(Store, optional): Key/value store accessible to all workers, used</span>
<span class="sd">                                to exchange connection/address information.</span>
<span class="sd">                                Mutually exclusive with ``init_method``.</span>
<span class="sd">        timeout (timedelta, optional): Timeout for operations executed against</span>
<span class="sd">            the process group. Default value equals 30 minutes.</span>
<span class="sd">            This is applicable for the ``gloo`` backend. For ``nccl``, this is</span>
<span class="sd">            applicable only if the environment variable ``NCCL_BLOCKING_WAIT``</span>
<span class="sd">            is set to 1.</span>
<span class="sd">        group_name (str, optional, deprecated): Group name.</span>

<span class="sd">    To enable ``backend == Backend.MPI``, PyTorch needs to built from source</span>
<span class="sd">    on a system that supports MPI.</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">global</span> <span class="n">_pg_group_ranks</span>
    <span class="k">global</span> <span class="n">_backend</span>
    <span class="k">global</span> <span class="n">_default_pg</span>
    <span class="k">global</span> <span class="n">_default_pg_init_method</span>

    <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">timeout</span><span class="p">,</span> <span class="n">timedelta</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Expected timeout argument to be of type&quot;</span>
                           <span class="s2">&quot;datetime.timedelta&quot;</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">_default_pg</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;trying to initialize the default process group &quot;</span>
                           <span class="s2">&quot;twice!&quot;</span><span class="p">)</span>

    <span class="k">assert</span> <span class="p">(</span><span class="n">store</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">)</span> <span class="ow">or</span> <span class="p">(</span><span class="n">init_method</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">),</span> \
        <span class="s2">&quot;Cannot specify both init_method and store.&quot;</span>

    <span class="k">if</span> <span class="n">store</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">assert</span> <span class="n">world_size</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">,</span> <span class="s1">&#39;world_size must be positive if using store&#39;</span>
        <span class="k">assert</span> <span class="n">rank</span> <span class="o">&gt;=</span> <span class="mi">0</span><span class="p">,</span> <span class="s1">&#39;rank must be non-negative if using store&#39;</span>
    <span class="k">elif</span> <span class="n">init_method</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">init_method</span> <span class="o">=</span> <span class="s2">&quot;env://&quot;</span>

    <span class="n">backend</span> <span class="o">=</span> <span class="n">Backend</span><span class="p">(</span><span class="n">backend</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">backend</span> <span class="o">==</span> <span class="n">Backend</span><span class="o">.</span><span class="n">MPI</span><span class="p">:</span>
        <span class="n">_default_pg</span> <span class="o">=</span> <span class="n">_new_process_group_helper</span><span class="p">(</span>
            <span class="o">-</span><span class="mi">1</span><span class="p">,</span>
            <span class="o">-</span><span class="mi">1</span><span class="p">,</span>
            <span class="p">[],</span>
            <span class="n">Backend</span><span class="o">.</span><span class="n">MPI</span><span class="p">,</span>
            <span class="kc">None</span><span class="p">,</span>
            <span class="n">group_name</span><span class="o">=</span><span class="n">group_name</span><span class="p">,</span>
            <span class="n">timeout</span><span class="o">=</span><span class="n">timeout</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="c1"># backward compatible API</span>
        <span class="k">if</span> <span class="n">store</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">rendezvous_iterator</span> <span class="o">=</span> <span class="n">rendezvous</span><span class="p">(</span>
                <span class="n">init_method</span><span class="p">,</span> <span class="n">rank</span><span class="p">,</span> <span class="n">world_size</span><span class="p">,</span> <span class="n">timeout</span><span class="o">=</span><span class="n">timeout</span>
            <span class="p">)</span>
            <span class="n">store</span><span class="p">,</span> <span class="n">rank</span><span class="p">,</span> <span class="n">world_size</span> <span class="o">=</span> <span class="nb">next</span><span class="p">(</span><span class="n">rendezvous_iterator</span><span class="p">)</span>
            <span class="n">store</span><span class="o">.</span><span class="n">set_timeout</span><span class="p">(</span><span class="n">timeout</span><span class="p">)</span>

        <span class="n">_default_pg</span> <span class="o">=</span> <span class="n">_new_process_group_helper</span><span class="p">(</span>
            <span class="n">world_size</span><span class="p">,</span>
            <span class="n">rank</span><span class="p">,</span>
            <span class="p">[],</span>
            <span class="n">backend</span><span class="p">,</span>
            <span class="n">store</span><span class="p">,</span>
            <span class="n">group_name</span><span class="o">=</span><span class="n">group_name</span><span class="p">,</span>
            <span class="n">timeout</span><span class="o">=</span><span class="n">timeout</span><span class="p">)</span>

    <span class="n">_pg_group_ranks</span><span class="p">[</span><span class="n">_default_pg</span><span class="p">]</span> <span class="o">=</span> <span class="p">{</span><span class="n">i</span><span class="p">:</span> <span class="n">i</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">_default_pg</span><span class="o">.</span><span class="n">size</span><span class="p">())}</span>
    <span class="n">_backend</span> <span class="o">=</span> <span class="n">_pg_map</span><span class="p">[</span><span class="n">_default_pg</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span>
    <span class="n">_default_pg_init_method</span> <span class="o">=</span> <span class="n">init_method</span></div>


<span class="k">def</span> <span class="nf">_new_process_group_helper</span><span class="p">(</span><span class="n">world_size</span><span class="p">,</span>
                              <span class="n">rank</span><span class="p">,</span>
                              <span class="n">group_ranks</span><span class="p">,</span>
                              <span class="n">backend</span><span class="p">,</span>
                              <span class="n">store</span><span class="p">,</span>
                              <span class="n">group_name</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
                              <span class="n">timeout</span><span class="o">=</span><span class="n">default_pg_timeout</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Create a new distributed process group.</span>

<span class="sd">    This function must be called by ALL processes in the global group, even if</span>
<span class="sd">    the calling process is not part of the newly created group. In that case,</span>
<span class="sd">    this function returns GroupMember.NON_GROUP_MEMBER.</span>

<span class="sd">    This function is called with ``group_ranks == []`` for the default group.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">global</span> <span class="n">_pg_map</span>
    <span class="k">global</span> <span class="n">_group_count</span>
    <span class="k">global</span> <span class="n">_pg_names</span>

    <span class="k">if</span> <span class="ow">not</span> <span class="n">group_name</span><span class="p">:</span>
        <span class="n">group_name</span> <span class="o">=</span> <span class="nb">str</span><span class="p">(</span><span class="n">_group_count</span><span class="p">)</span>
        <span class="n">_group_count</span> <span class="o">+=</span> <span class="mi">1</span>

    <span class="k">if</span> <span class="n">group_name</span> <span class="ow">in</span> <span class="n">_pg_names</span><span class="o">.</span><span class="n">values</span><span class="p">():</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;The specified group name has already been &quot;</span>
                           <span class="s2">&quot;created, please use a different group name&quot;</span><span class="p">)</span>

    <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">timeout</span><span class="p">,</span> <span class="n">timedelta</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Expected timeout argument to be of type&quot;</span>
                           <span class="s2">&quot;datetime.timedelta&quot;</span><span class="p">)</span>

    <span class="c1"># The list of group ranks is empty if we&#39;re creating the default group.</span>
    <span class="n">is_default_group</span> <span class="o">=</span> <span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">group_ranks</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">)</span>

    <span class="n">backend</span> <span class="o">=</span> <span class="n">Backend</span><span class="p">(</span><span class="n">backend</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">backend</span> <span class="o">==</span> <span class="n">Backend</span><span class="o">.</span><span class="n">MPI</span><span class="p">:</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">is_mpi_available</span><span class="p">():</span>
            <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Distributed package doesn&#39;t have MPI built in&quot;</span><span class="p">)</span>
        <span class="n">pg</span> <span class="o">=</span> <span class="n">ProcessGroupMPI</span><span class="o">.</span><span class="n">create</span><span class="p">(</span><span class="n">group_ranks</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">pg</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">NON_GROUP_MEMBER</span>
        <span class="n">_pg_map</span><span class="p">[</span><span class="n">pg</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span><span class="n">Backend</span><span class="o">.</span><span class="n">MPI</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
        <span class="n">_pg_names</span><span class="p">[</span><span class="n">pg</span><span class="p">]</span> <span class="o">=</span> <span class="n">group_name</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="c1"># If this is a subgroup (which means group_ranks is specified),</span>
        <span class="c1"># we check if the current process is a member of the new group.</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">is_default_group</span><span class="p">:</span>
            <span class="n">global_rank</span> <span class="o">=</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">rank</span><span class="p">()</span>
            <span class="k">if</span> <span class="n">global_rank</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">group_ranks</span><span class="p">:</span>
                <span class="k">return</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">NON_GROUP_MEMBER</span>

        <span class="c1"># Use the group name as prefix in the default store, such that</span>
        <span class="c1"># a single store can be reused by multiple groups.</span>
        <span class="n">prefix_store</span> <span class="o">=</span> <span class="n">PrefixStore</span><span class="p">(</span><span class="n">group_name</span><span class="p">,</span> <span class="n">store</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">backend</span> <span class="o">==</span> <span class="n">Backend</span><span class="o">.</span><span class="n">GLOO</span><span class="p">:</span>
            <span class="n">pg</span> <span class="o">=</span> <span class="n">ProcessGroupGloo</span><span class="p">(</span>
                <span class="n">prefix_store</span><span class="p">,</span>
                <span class="n">rank</span><span class="p">,</span>
                <span class="n">world_size</span><span class="p">,</span>
                <span class="n">timeout</span><span class="o">=</span><span class="n">timeout</span><span class="p">)</span>
            <span class="n">_pg_map</span><span class="p">[</span><span class="n">pg</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span><span class="n">Backend</span><span class="o">.</span><span class="n">GLOO</span><span class="p">,</span> <span class="n">store</span><span class="p">)</span>
            <span class="n">_pg_names</span><span class="p">[</span><span class="n">pg</span><span class="p">]</span> <span class="o">=</span> <span class="n">group_name</span>
        <span class="k">elif</span> <span class="n">backend</span> <span class="o">==</span> <span class="n">Backend</span><span class="o">.</span><span class="n">NCCL</span><span class="p">:</span>
            <span class="k">if</span> <span class="ow">not</span> <span class="n">is_nccl_available</span><span class="p">():</span>
                <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Distributed package doesn&#39;t have NCCL &quot;</span>
                                   <span class="s2">&quot;built in&quot;</span><span class="p">)</span>
            <span class="n">pg</span> <span class="o">=</span> <span class="n">ProcessGroupNCCL</span><span class="p">(</span>
                <span class="n">prefix_store</span><span class="p">,</span>
                <span class="n">rank</span><span class="p">,</span>
                <span class="n">world_size</span><span class="p">,</span>
                <span class="n">timeout</span><span class="p">)</span>
            <span class="n">_pg_map</span><span class="p">[</span><span class="n">pg</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span><span class="n">Backend</span><span class="o">.</span><span class="n">NCCL</span><span class="p">,</span> <span class="n">store</span><span class="p">)</span>
            <span class="n">_pg_names</span><span class="p">[</span><span class="n">pg</span><span class="p">]</span> <span class="o">=</span> <span class="n">group_name</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Unsupported distributed backend by group&quot;</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">pg</span>


<span class="k">def</span> <span class="nf">destroy_process_group</span><span class="p">(</span><span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Destroy a given process group, and deinitialize the distributed package</span>

<span class="sd">    Arguments:</span>
<span class="sd">        group (ProcessGroup, optional): The process group to be destroyed, if</span>
<span class="sd">                                        group.WORLD is given, all process</span>
<span class="sd">                                        groups including the default one will</span>
<span class="sd">                                        be destroyed.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">global</span> <span class="n">_pg_map</span>
    <span class="k">global</span> <span class="n">_pg_names</span>
    <span class="k">global</span> <span class="n">_pg_group_ranks</span>
    <span class="k">global</span> <span class="n">_default_pg</span>
    <span class="k">global</span> <span class="n">_default_pg_init_method</span>
    <span class="k">global</span> <span class="n">_group_count</span>

    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">NON_GROUP_MEMBER</span><span class="p">:</span>
        <span class="k">return</span>

    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">pg</span> <span class="o">=</span> <span class="n">_default_pg</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">pg</span> <span class="o">=</span> <span class="n">group</span>

    <span class="k">if</span> <span class="n">_pg_map</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">pg</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Invalid process group specified&quot;</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_default_pg</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="n">_default_pg_init_method</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="n">_pg_map</span><span class="o">.</span><span class="n">clear</span><span class="p">()</span>
        <span class="n">_pg_names</span><span class="o">.</span><span class="n">clear</span><span class="p">()</span>
        <span class="n">_pg_group_ranks</span><span class="o">.</span><span class="n">clear</span><span class="p">()</span>

        <span class="c1"># when process group doesn&#39;t have an explicit name (only WORLD (default)</span>
        <span class="c1"># process group can have an explicit name), we use global _group_counter</span>
        <span class="c1"># to generate the name. We need to reset the counter on destruction to</span>
        <span class="c1"># allow consistent value to be generated when we re-create process</span>
        <span class="c1"># groups after some trainers recover from failure</span>
        <span class="c1">#</span>
        <span class="c1"># We only reset this when WORLD is being destroyed because if this</span>
        <span class="c1"># process group is in good state, we aren&#39;t dealing with failures.</span>
        <span class="n">_group_count</span> <span class="o">=</span> <span class="mi">0</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">del</span> <span class="n">_pg_map</span><span class="p">[</span><span class="n">pg</span><span class="p">]</span>
        <span class="k">del</span> <span class="n">_pg_names</span><span class="p">[</span><span class="n">pg</span><span class="p">]</span>
        <span class="k">del</span> <span class="n">_pg_group_ranks</span><span class="p">[</span><span class="n">pg</span><span class="p">]</span>


<div class="viewcode-block" id="get_rank"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.get_rank">[docs]</a><span class="k">def</span> <span class="nf">get_rank</span><span class="p">(</span><span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Returns the rank of current process group</span>

<span class="sd">    Rank is a unique identifier assigned to each process within a distributed</span>
<span class="sd">    process group. They are always consecutive integers ranging from 0 to</span>
<span class="sd">    ``world_size``.</span>

<span class="sd">    Arguments:</span>
<span class="sd">        group (ProcessGroup, optional): The process group to work on</span>

<span class="sd">    Returns:</span>
<span class="sd">        The rank of the process group</span>
<span class="sd">        -1, if not part of the group</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span> <span class="o">-</span><span class="mi">1</span>

    <span class="n">_check_default_pg</span><span class="p">()</span>
    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">rank</span><span class="p">()</span>

    <span class="k">return</span> <span class="n">_get_group_rank</span><span class="p">(</span><span class="n">group</span><span class="p">,</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">rank</span><span class="p">())</span></div>


<div class="viewcode-block" id="get_world_size"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.get_world_size">[docs]</a><span class="k">def</span> <span class="nf">get_world_size</span><span class="p">(</span><span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Returns the number of processes in the current process group</span>

<span class="sd">    Arguments:</span>
<span class="sd">        group (ProcessGroup, optional): The process group to work on</span>

<span class="sd">    Returns:</span>
<span class="sd">        The world size of the process group</span>
<span class="sd">        -1, if not part of the group</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span> <span class="o">-</span><span class="mi">1</span>

    <span class="k">return</span> <span class="n">_get_group_size</span><span class="p">(</span><span class="n">group</span><span class="p">)</span></div>


<div class="viewcode-block" id="isend"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.isend">[docs]</a><span class="k">def</span> <span class="nf">isend</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span>
          <span class="n">dst</span><span class="p">,</span>
          <span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">,</span>
          <span class="n">tag</span><span class="o">=</span><span class="mi">0</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Sends a tensor asynchronously.</span>

<span class="sd">    Arguments:</span>
<span class="sd">        tensor (Tensor): Tensor to send.</span>
<span class="sd">        dst (int): Destination rank.</span>
<span class="sd">        group (ProcessGroup, optional): The process group to work on</span>
<span class="sd">        tag (int, optional): Tag to match send with remote recv</span>

<span class="sd">    Returns:</span>
<span class="sd">        A distributed request object.</span>
<span class="sd">        None, if not part of the group</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">_check_single_tensor</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="s2">&quot;tensor&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span>

    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_check_default_pg</span><span class="p">()</span>
        <span class="k">return</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">send</span><span class="p">([</span><span class="n">tensor</span><span class="p">],</span> <span class="n">dst</span><span class="p">,</span> <span class="n">tag</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">group_dst_rank</span> <span class="o">=</span> <span class="n">_get_group_rank</span><span class="p">(</span><span class="n">group</span><span class="p">,</span> <span class="n">dst</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">group</span><span class="o">.</span><span class="n">send</span><span class="p">([</span><span class="n">tensor</span><span class="p">],</span> <span class="n">group_dst_rank</span><span class="p">,</span> <span class="n">tag</span><span class="p">)</span></div>


<div class="viewcode-block" id="irecv"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.irecv">[docs]</a><span class="k">def</span> <span class="nf">irecv</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span>
          <span class="n">src</span><span class="p">,</span>
          <span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">,</span>
          <span class="n">tag</span><span class="o">=</span><span class="mi">0</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Receives a tensor asynchronously.</span>

<span class="sd">    Arguments:</span>
<span class="sd">        tensor (Tensor): Tensor to fill with received data.</span>
<span class="sd">        src (int): Source rank.</span>
<span class="sd">        group (ProcessGroup, optional): The process group to work on</span>
<span class="sd">        tag (int, optional): Tag to match recv with remote send</span>

<span class="sd">    Returns:</span>
<span class="sd">        A distributed request object.</span>
<span class="sd">        None, if not part of the group</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">_check_single_tensor</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="s2">&quot;tensor&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span>

    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_check_default_pg</span><span class="p">()</span>
        <span class="k">return</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">recv</span><span class="p">([</span><span class="n">tensor</span><span class="p">],</span> <span class="n">src</span><span class="p">,</span> <span class="n">tag</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">group_src_rank</span> <span class="o">=</span> <span class="n">_get_group_rank</span><span class="p">(</span><span class="n">group</span><span class="p">,</span> <span class="n">src</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">group</span><span class="o">.</span><span class="n">recv</span><span class="p">([</span><span class="n">tensor</span><span class="p">],</span> <span class="n">group_src_rank</span><span class="p">,</span> <span class="n">tag</span><span class="p">)</span></div>


<div class="viewcode-block" id="send"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.send">[docs]</a><span class="k">def</span> <span class="nf">send</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span>
         <span class="n">dst</span><span class="p">,</span>
         <span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">,</span>
         <span class="n">tag</span><span class="o">=</span><span class="mi">0</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Sends a tensor synchronously.</span>

<span class="sd">    Arguments:</span>
<span class="sd">        tensor (Tensor): Tensor to send.</span>
<span class="sd">        dst (int): Destination rank.</span>
<span class="sd">        group (ProcessGroup, optional): The process group to work on</span>
<span class="sd">        tag (int, optional): Tag to match send with remote recv</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">_check_single_tensor</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="s2">&quot;tensor&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span>

    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_check_default_pg</span><span class="p">()</span>
        <span class="n">_default_pg</span><span class="o">.</span><span class="n">send</span><span class="p">([</span><span class="n">tensor</span><span class="p">],</span> <span class="n">dst</span><span class="p">,</span> <span class="n">tag</span><span class="p">)</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">group_dst_rank</span> <span class="o">=</span> <span class="n">_get_group_rank</span><span class="p">(</span><span class="n">group</span><span class="p">,</span> <span class="n">dst</span><span class="p">)</span>
        <span class="n">group</span><span class="o">.</span><span class="n">send</span><span class="p">([</span><span class="n">tensor</span><span class="p">],</span> <span class="n">group_dst_rank</span><span class="p">,</span> <span class="n">tag</span><span class="p">)</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span></div>


<div class="viewcode-block" id="recv"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.recv">[docs]</a><span class="k">def</span> <span class="nf">recv</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span>
         <span class="n">src</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
         <span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">,</span>
         <span class="n">tag</span><span class="o">=</span><span class="mi">0</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Receives a tensor synchronously.</span>

<span class="sd">    Arguments:</span>
<span class="sd">        tensor (Tensor): Tensor to fill with received data.</span>
<span class="sd">        src (int, optional): Source rank. Will receive from any</span>
<span class="sd">            process if unspecified.</span>
<span class="sd">        group (ProcessGroup, optional): The process group to work on</span>
<span class="sd">        tag (int, optional): Tag to match recv with remote send</span>

<span class="sd">    Returns:</span>
<span class="sd">        Sender rank</span>
<span class="sd">        -1, if not part of the group</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">_check_single_tensor</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="s2">&quot;tensor&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span> <span class="o">-</span><span class="mi">1</span>

    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_check_default_pg</span><span class="p">()</span>
        <span class="n">pg</span> <span class="o">=</span> <span class="n">_default_pg</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">pg</span> <span class="o">=</span> <span class="n">group</span>

    <span class="k">if</span> <span class="n">src</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">pg</span><span class="o">.</span><span class="n">recv_anysource</span><span class="p">([</span><span class="n">tensor</span><span class="p">],</span> <span class="n">tag</span><span class="p">)</span>
        <span class="n">work</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span>
        <span class="n">src_rank</span> <span class="o">=</span> <span class="n">work</span><span class="o">.</span><span class="n">source_rank</span><span class="p">()</span>
        <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">src_rank</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">return</span> <span class="n">_get_global_rank</span><span class="p">(</span><span class="n">pg</span><span class="p">,</span> <span class="n">src_rank</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
            <span class="n">pg</span><span class="o">.</span><span class="n">recv</span><span class="p">([</span><span class="n">tensor</span><span class="p">],</span> <span class="n">src</span><span class="p">,</span> <span class="n">tag</span><span class="p">)</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">group_src_rank</span> <span class="o">=</span> <span class="n">_get_group_rank</span><span class="p">(</span><span class="n">pg</span><span class="p">,</span> <span class="n">src</span><span class="p">)</span>
            <span class="n">pg</span><span class="o">.</span><span class="n">recv</span><span class="p">([</span><span class="n">tensor</span><span class="p">],</span> <span class="n">group_src_rank</span><span class="p">,</span> <span class="n">tag</span><span class="p">)</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span>
        <span class="k">return</span> <span class="n">src</span></div>


<div class="viewcode-block" id="broadcast_multigpu"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.broadcast_multigpu">[docs]</a><span class="k">def</span> <span class="nf">broadcast_multigpu</span><span class="p">(</span><span class="n">tensor_list</span><span class="p">,</span>
                       <span class="n">src</span><span class="p">,</span>
                       <span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">,</span>
                       <span class="n">async_op</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                       <span class="n">src_tensor</span><span class="o">=</span><span class="mi">0</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Broadcasts the tensor to the whole group with multiple GPU tensors</span>
<span class="sd">    per node.</span>

<span class="sd">    ``tensor`` must have the same number of elements in all the GPUs from</span>
<span class="sd">    all processes participating in the collective. each tensor in the list must</span>
<span class="sd">    be on a different GPU</span>

<span class="sd">    Only nccl and gloo backend are currently supported</span>
<span class="sd">    tensors should only be GPU tensors</span>

<span class="sd">    Arguments:</span>
<span class="sd">        tensor_list (List[Tensor]): Tensors that participate in the collective</span>
<span class="sd">            operation. If ``src`` is the rank, then the specified ``src_tensor``</span>
<span class="sd">            element of ``tensor_list`` (``tensor_list[src_tensor]``) will be</span>
<span class="sd">            broadcast to all other tensors (on different GPUs) in the src process</span>
<span class="sd">            and all tensors in ``tensor_list`` of other non-src processes.</span>
<span class="sd">            You also need to make sure that ``len(tensor_list)`` is the same</span>
<span class="sd">            for all the distributed processes calling this function.</span>

<span class="sd">        src (int): Source rank.</span>
<span class="sd">        group (ProcessGroup, optional): The process group to work on</span>
<span class="sd">        async_op (bool, optional): Whether this op should be an async op</span>
<span class="sd">        src_tensor (int, optional): Source tensor rank within ``tensor_list``</span>

<span class="sd">    Returns:</span>
<span class="sd">        Async work handle, if async_op is set to True.</span>
<span class="sd">        None, if not async_op or if not part of the group</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span>

    <span class="n">opts</span> <span class="o">=</span> <span class="n">BroadcastOptions</span><span class="p">()</span>
    <span class="n">opts</span><span class="o">.</span><span class="n">rootRank</span> <span class="o">=</span> <span class="n">src</span>
    <span class="n">opts</span><span class="o">.</span><span class="n">rootTensor</span> <span class="o">=</span> <span class="n">src_tensor</span>

    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_check_default_pg</span><span class="p">()</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">broadcast</span><span class="p">(</span><span class="n">tensor_list</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">group_src_rank</span> <span class="o">=</span> <span class="n">_get_group_rank</span><span class="p">(</span><span class="n">group</span><span class="p">,</span> <span class="n">src</span><span class="p">)</span>
        <span class="n">opts</span><span class="o">.</span><span class="n">rootRank</span> <span class="o">=</span> <span class="n">group_src_rank</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">group</span><span class="o">.</span><span class="n">broadcast</span><span class="p">(</span><span class="n">tensor_list</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">async_op</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">work</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span></div>


<div class="viewcode-block" id="broadcast"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.broadcast">[docs]</a><span class="k">def</span> <span class="nf">broadcast</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span>
              <span class="n">src</span><span class="p">,</span>
              <span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">,</span>
              <span class="n">async_op</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Broadcasts the tensor to the whole group.</span>

<span class="sd">    ``tensor`` must have the same number of elements in all processes</span>
<span class="sd">    participating in the collective.</span>

<span class="sd">    Arguments:</span>
<span class="sd">        tensor (Tensor): Data to be sent if ``src`` is the rank of current</span>
<span class="sd">            process, and tensor to be used to save received data otherwise.</span>
<span class="sd">        src (int): Source rank.</span>
<span class="sd">        group (ProcessGroup, optional): The process group to work on</span>
<span class="sd">        async_op (bool, optional): Whether this op should be an async op</span>

<span class="sd">    Returns:</span>
<span class="sd">        Async work handle, if async_op is set to True.</span>
<span class="sd">        None, if not async_op or if not part of the group</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">_check_single_tensor</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="s2">&quot;tensor&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span>

    <span class="n">opts</span> <span class="o">=</span> <span class="n">BroadcastOptions</span><span class="p">()</span>
    <span class="n">opts</span><span class="o">.</span><span class="n">rootRank</span> <span class="o">=</span> <span class="n">src</span>
    <span class="n">opts</span><span class="o">.</span><span class="n">rootTensor</span> <span class="o">=</span> <span class="mi">0</span>

    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_check_default_pg</span><span class="p">()</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">broadcast</span><span class="p">([</span><span class="n">tensor</span><span class="p">],</span> <span class="n">opts</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">group_src_rank</span> <span class="o">=</span> <span class="n">_get_group_rank</span><span class="p">(</span><span class="n">group</span><span class="p">,</span> <span class="n">src</span><span class="p">)</span>
        <span class="n">opts</span><span class="o">.</span><span class="n">rootRank</span> <span class="o">=</span> <span class="n">group_src_rank</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">group</span><span class="o">.</span><span class="n">broadcast</span><span class="p">([</span><span class="n">tensor</span><span class="p">],</span> <span class="n">opts</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">async_op</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">work</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span></div>


<div class="viewcode-block" id="all_reduce_multigpu"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.all_reduce_multigpu">[docs]</a><span class="k">def</span> <span class="nf">all_reduce_multigpu</span><span class="p">(</span><span class="n">tensor_list</span><span class="p">,</span>
                        <span class="n">op</span><span class="o">=</span><span class="n">ReduceOp</span><span class="o">.</span><span class="n">SUM</span><span class="p">,</span>
                        <span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">,</span>
                        <span class="n">async_op</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Reduces the tensor data across all machines in such a way that all get</span>
<span class="sd">    the final result. This function reduces a number of tensors on every node,</span>
<span class="sd">    while each tensor resides on different GPUs.</span>
<span class="sd">    Therefore, the input tensor in the tensor list needs to be GPU tensors.</span>
<span class="sd">    Also, each tensor in the tensor list needs to reside on a different GPU.</span>

<span class="sd">    After the call, all ``tensor`` in ``tensor_list`` is going to be bitwise</span>
<span class="sd">    identical in all processes.</span>

<span class="sd">    Only nccl and gloo backend is currently supported</span>
<span class="sd">    tensors should only be GPU tensors</span>

<span class="sd">    Arguments:</span>
<span class="sd">        tensor list (List[Tensor]): List of input and output tensors of</span>
<span class="sd">            the collective. The function operates in-place and requires that</span>
<span class="sd">            each tensor to be a GPU tensor on different GPUs.</span>
<span class="sd">            You also need to make sure that ``len(tensor_list)`` is the same for</span>
<span class="sd">            all the distributed processes calling this function.</span>
<span class="sd">        op (optional): One of the values from</span>
<span class="sd">            ``torch.distributed.ReduceOp``</span>
<span class="sd">            enum.  Specifies an operation used for element-wise reductions.</span>
<span class="sd">        group (ProcessGroup, optional): The process group to work on</span>
<span class="sd">        async_op (bool, optional): Whether this op should be an async op</span>

<span class="sd">    Returns:</span>
<span class="sd">        Async work handle, if async_op is set to True.</span>
<span class="sd">        None, if not async_op or if not part of the group</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span>

    <span class="n">opts</span> <span class="o">=</span> <span class="n">AllreduceOptions</span><span class="p">()</span>
    <span class="n">opts</span><span class="o">.</span><span class="n">reduceOp</span> <span class="o">=</span> <span class="n">op</span>
    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_check_default_pg</span><span class="p">()</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">allreduce</span><span class="p">(</span><span class="n">tensor_list</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">group</span><span class="o">.</span><span class="n">allreduce</span><span class="p">(</span><span class="n">tensor_list</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">async_op</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">work</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span></div>


<div class="viewcode-block" id="all_reduce"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.all_reduce">[docs]</a><span class="k">def</span> <span class="nf">all_reduce</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span>
               <span class="n">op</span><span class="o">=</span><span class="n">ReduceOp</span><span class="o">.</span><span class="n">SUM</span><span class="p">,</span>
               <span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">,</span>
               <span class="n">async_op</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Reduces the tensor data across all machines in such a way that all get</span>
<span class="sd">    the final result.</span>

<span class="sd">    After the call ``tensor`` is going to be bitwise identical in all processes.</span>

<span class="sd">    Arguments:</span>
<span class="sd">        tensor (Tensor): Input and output of the collective. The function</span>
<span class="sd">            operates in-place.</span>
<span class="sd">        op (optional): One of the values from</span>
<span class="sd">            ``torch.distributed.ReduceOp``</span>
<span class="sd">            enum.  Specifies an operation used for element-wise reductions.</span>
<span class="sd">        group (ProcessGroup, optional): The process group to work on</span>
<span class="sd">        async_op (bool, optional): Whether this op should be an async op</span>

<span class="sd">    Returns:</span>
<span class="sd">        Async work handle, if async_op is set to True.</span>
<span class="sd">        None, if not async_op or if not part of the group</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">_check_single_tensor</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="s2">&quot;tensor&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span>

    <span class="n">opts</span> <span class="o">=</span> <span class="n">AllreduceOptions</span><span class="p">()</span>
    <span class="n">opts</span><span class="o">.</span><span class="n">reduceOp</span> <span class="o">=</span> <span class="n">op</span>
    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_check_default_pg</span><span class="p">()</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">allreduce</span><span class="p">([</span><span class="n">tensor</span><span class="p">],</span> <span class="n">opts</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">group</span><span class="o">.</span><span class="n">allreduce</span><span class="p">([</span><span class="n">tensor</span><span class="p">],</span> <span class="n">opts</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">async_op</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">work</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span></div>


<span class="k">def</span> <span class="nf">all_reduce_coalesced</span><span class="p">(</span><span class="n">tensors</span><span class="p">,</span>
                         <span class="n">op</span><span class="o">=</span><span class="n">ReduceOp</span><span class="o">.</span><span class="n">SUM</span><span class="p">,</span>
                         <span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">,</span>
                         <span class="n">async_op</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    WARNING: at this time individual shape checking is not implemented across nodes.</span>
<span class="sd">    For example, if the rank 0 node passes [torch.rand(4), torch.rand(2)] and the</span>
<span class="sd">    rank 1 node passes [torch.rand(2), torch.rand(2), torch.rand(2)], the allreduce</span>
<span class="sd">    operation will proceed without complaint and return erroneous outputs. This lack</span>
<span class="sd">    of shape checking results in significant performance improvements but users of this</span>
<span class="sd">    function should take extra care to ensure that each node passes in tensors whose</span>
<span class="sd">    shapes match across nodes.</span>

<span class="sd">    Reduces each tensor in tensors (residing on the same device) across all machines</span>
<span class="sd">    in such a way that all get the final result.</span>

<span class="sd">    After the call each tensor in tensors is going to bitwise identical</span>
<span class="sd">    in all processes.</span>

<span class="sd">    Arguments:</span>
<span class="sd">        tensors (List[Tensor]): Input and output of the collective. The function</span>
<span class="sd">            operates in-place.</span>
<span class="sd">        op (Optional[ReduceOp]): One of the values from</span>
<span class="sd">            ``torch.distributed.ReduceOp`` enum. Specifies an operation used for</span>
<span class="sd">            element-wise reductions.</span>
<span class="sd">        group (Optional[ProcessGroup]): The process group to work on.</span>
<span class="sd">        async_op (Optional[bool]): Whether this op should be an async op.</span>

<span class="sd">    Returns:</span>
<span class="sd">        Async work handle, if async_op is set to True.</span>
<span class="sd">        None, if not async_op or if not part of the group.</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">_check_tensor_list</span><span class="p">(</span><span class="n">tensors</span><span class="p">,</span> <span class="s2">&quot;tensor&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span>

    <span class="n">opts</span> <span class="o">=</span> <span class="n">AllreduceCoalescedOptions</span><span class="p">()</span>
    <span class="n">opts</span><span class="o">.</span><span class="n">reduceOp</span> <span class="o">=</span> <span class="n">op</span>
    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_check_default_pg</span><span class="p">()</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">allreduce_coalesced</span><span class="p">(</span><span class="n">tensors</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">group</span><span class="o">.</span><span class="n">allreduce_coalesced</span><span class="p">(</span><span class="n">tensors</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">async_op</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">work</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span>


<div class="viewcode-block" id="reduce_multigpu"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.reduce_multigpu">[docs]</a><span class="k">def</span> <span class="nf">reduce_multigpu</span><span class="p">(</span><span class="n">tensor_list</span><span class="p">,</span>
                    <span class="n">dst</span><span class="p">,</span>
                    <span class="n">op</span><span class="o">=</span><span class="n">ReduceOp</span><span class="o">.</span><span class="n">SUM</span><span class="p">,</span>
                    <span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">,</span>
                    <span class="n">async_op</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                    <span class="n">dst_tensor</span><span class="o">=</span><span class="mi">0</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Reduces the tensor data on multiple GPUs across all machines. Each tensor</span>
<span class="sd">    in ``tensor_list`` should reside on a separate GPU</span>

<span class="sd">    Only the GPU of ``tensor_list[dst_tensor]`` on the process with rank ``dst``</span>
<span class="sd">    is going to receive the final result.</span>

<span class="sd">    Only nccl backend is currently supported</span>
<span class="sd">    tensors should only be GPU tensors</span>

<span class="sd">    Arguments:</span>
<span class="sd">        tensor_list (List[Tensor]): Input and output GPU tensors of the</span>
<span class="sd">            collective. The function operates in-place.</span>
<span class="sd">            You also need to make sure that ``len(tensor_list)`` is the same for</span>
<span class="sd">            all the distributed processes calling this function.</span>
<span class="sd">        dst (int): Destination rank</span>
<span class="sd">        op (optional): One of the values from</span>
<span class="sd">            ``torch.distributed.ReduceOp``</span>
<span class="sd">            enum.  Specifies an operation used for element-wise reductions.</span>
<span class="sd">        group (ProcessGroup, optional): The process group to work on</span>
<span class="sd">        async_op (bool, optional): Whether this op should be an async op</span>
<span class="sd">        dst_tensor (int, optional): Destination tensor rank within</span>
<span class="sd">                                    ``tensor_list``</span>

<span class="sd">    Returns:</span>
<span class="sd">        Async work handle, if async_op is set to True.</span>
<span class="sd">        None, otherwise</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span>

    <span class="n">opts</span> <span class="o">=</span> <span class="n">ReduceOptions</span><span class="p">()</span>
    <span class="n">opts</span><span class="o">.</span><span class="n">reduceOp</span> <span class="o">=</span> <span class="n">op</span>
    <span class="n">opts</span><span class="o">.</span><span class="n">rootRank</span> <span class="o">=</span> <span class="n">dst</span>
    <span class="n">opts</span><span class="o">.</span><span class="n">rootTensor</span> <span class="o">=</span> <span class="n">dst_tensor</span>

    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_check_default_pg</span><span class="p">()</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="n">tensor_list</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">group_dst_rank</span> <span class="o">=</span> <span class="n">_get_group_rank</span><span class="p">(</span><span class="n">group</span><span class="p">,</span> <span class="n">dst</span><span class="p">)</span>
        <span class="n">opts</span><span class="o">.</span><span class="n">rootRank</span> <span class="o">=</span> <span class="n">group_dst_rank</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">group</span><span class="o">.</span><span class="n">reduce</span><span class="p">(</span><span class="n">tensor_list</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">async_op</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">work</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span></div>


<div class="viewcode-block" id="reduce"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.reduce">[docs]</a><span class="k">def</span> <span class="nf">reduce</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span>
           <span class="n">dst</span><span class="p">,</span>
           <span class="n">op</span><span class="o">=</span><span class="n">ReduceOp</span><span class="o">.</span><span class="n">SUM</span><span class="p">,</span>
           <span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">,</span>
           <span class="n">async_op</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Reduces the tensor data across all machines.</span>

<span class="sd">    Only the process with rank ``dst`` is going to receive the final result.</span>

<span class="sd">    Arguments:</span>
<span class="sd">        tensor (Tensor): Input and output of the collective. The function</span>
<span class="sd">            operates in-place.</span>
<span class="sd">        dst (int): Destination rank</span>
<span class="sd">        op (optional): One of the values from</span>
<span class="sd">            ``torch.distributed.ReduceOp``</span>
<span class="sd">            enum.  Specifies an operation used for element-wise reductions.</span>
<span class="sd">        group (ProcessGroup, optional): The process group to work on</span>
<span class="sd">        async_op (bool, optional): Whether this op should be an async op</span>

<span class="sd">    Returns:</span>
<span class="sd">        Async work handle, if async_op is set to True.</span>
<span class="sd">        None, if not async_op or if not part of the group</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">_check_single_tensor</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="s2">&quot;tensor&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span>

    <span class="n">opts</span> <span class="o">=</span> <span class="n">ReduceOptions</span><span class="p">()</span>
    <span class="n">opts</span><span class="o">.</span><span class="n">reduceOp</span> <span class="o">=</span> <span class="n">op</span>
    <span class="n">opts</span><span class="o">.</span><span class="n">rootRank</span> <span class="o">=</span> <span class="n">dst</span>

    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_check_default_pg</span><span class="p">()</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">reduce</span><span class="p">([</span><span class="n">tensor</span><span class="p">],</span> <span class="n">opts</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">group_dst_rank</span> <span class="o">=</span> <span class="n">_get_group_rank</span><span class="p">(</span><span class="n">group</span><span class="p">,</span> <span class="n">dst</span><span class="p">)</span>
        <span class="n">opts</span><span class="o">.</span><span class="n">rootRank</span> <span class="o">=</span> <span class="n">group_dst_rank</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">group</span><span class="o">.</span><span class="n">reduce</span><span class="p">([</span><span class="n">tensor</span><span class="p">],</span> <span class="n">opts</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">async_op</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">work</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span></div>


<div class="viewcode-block" id="all_gather_multigpu"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.all_gather_multigpu">[docs]</a><span class="k">def</span> <span class="nf">all_gather_multigpu</span><span class="p">(</span><span class="n">output_tensor_lists</span><span class="p">,</span>
                        <span class="n">input_tensor_list</span><span class="p">,</span>
                        <span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">,</span>
                        <span class="n">async_op</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Gathers tensors from the whole group in a list.</span>
<span class="sd">    Each tensor in ``tensor_list`` should reside on a separate GPU</span>

<span class="sd">    Only nccl backend is currently supported</span>
<span class="sd">    tensors should only be GPU tensors</span>

<span class="sd">    Arguments:</span>
<span class="sd">        output_tensor_lists (List[List[Tensor]]): Output lists. It should</span>
<span class="sd">            contain correctly-sized tensors on each GPU to be used for output</span>
<span class="sd">            of the collective, e.g. ``output_tensor_lists[i]`` contains the</span>
<span class="sd">            all_gather result that resides on the GPU of</span>
<span class="sd">            ``input_tensor_list[i]``.</span>

<span class="sd">            Note that each element of ``output_tensor_lists`` has the size of</span>
<span class="sd">            ``world_size * len(input_tensor_list)``, since the function all</span>
<span class="sd">            gathers the result from every single GPU in the group. To interpret</span>
<span class="sd">            each element of ``output_tensor_lists[i]``, note that</span>
<span class="sd">            ``input_tensor_list[j]`` of rank k will be appear in</span>
<span class="sd">            ``output_tensor_lists[i][k * world_size + j]``</span>

<span class="sd">            Also note that ``len(output_tensor_lists)``, and the size of each</span>
<span class="sd">            element in ``output_tensor_lists`` (each element is a list,</span>
<span class="sd">            therefore ``len(output_tensor_lists[i])``) need to be the same</span>
<span class="sd">            for all the distributed processes calling this function.</span>

<span class="sd">        input_tensor_list (List[Tensor]): List of tensors(on different GPUs) to</span>
<span class="sd">            be broadcast from current process.</span>
<span class="sd">            Note that ``len(input_tensor_list)`` needs to be the same for</span>
<span class="sd">            all the distributed processes calling this function.</span>

<span class="sd">        group (ProcessGroup, optional): The process group to work on</span>
<span class="sd">        async_op (bool, optional): Whether this op should be an async op</span>

<span class="sd">    Returns:</span>
<span class="sd">        Async work handle, if async_op is set to True.</span>
<span class="sd">        None, if not async_op or if not part of the group</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span>

    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_check_default_pg</span><span class="p">()</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">allgather</span><span class="p">(</span><span class="n">output_tensor_lists</span><span class="p">,</span> <span class="n">input_tensor_list</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">group</span><span class="o">.</span><span class="n">allgather</span><span class="p">(</span><span class="n">output_tensor_lists</span><span class="p">,</span> <span class="n">input_tensor_list</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">async_op</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">work</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span></div>


<div class="viewcode-block" id="all_gather"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.all_gather">[docs]</a><span class="k">def</span> <span class="nf">all_gather</span><span class="p">(</span><span class="n">tensor_list</span><span class="p">,</span>
               <span class="n">tensor</span><span class="p">,</span>
               <span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">,</span>
               <span class="n">async_op</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Gathers tensors from the whole group in a list.</span>

<span class="sd">    Arguments:</span>
<span class="sd">        tensor_list (list[Tensor]): Output list. It should contain</span>
<span class="sd">            correctly-sized tensors to be used for output of the collective.</span>
<span class="sd">        tensor (Tensor): Tensor to be broadcast from current process.</span>
<span class="sd">        group (ProcessGroup, optional): The process group to work on</span>
<span class="sd">        async_op (bool, optional): Whether this op should be an async op</span>

<span class="sd">    Returns:</span>
<span class="sd">        Async work handle, if async_op is set to True.</span>
<span class="sd">        None, if not async_op or if not part of the group</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">_check_tensor_list</span><span class="p">(</span><span class="n">tensor_list</span><span class="p">,</span> <span class="s2">&quot;tensor_list&quot;</span><span class="p">)</span>
    <span class="n">_check_single_tensor</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="s2">&quot;tensor&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span>

    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_check_default_pg</span><span class="p">()</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">allgather</span><span class="p">([</span><span class="n">tensor_list</span><span class="p">],</span> <span class="p">[</span><span class="n">tensor</span><span class="p">])</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">group</span><span class="o">.</span><span class="n">allgather</span><span class="p">([</span><span class="n">tensor_list</span><span class="p">],</span> <span class="p">[</span><span class="n">tensor</span><span class="p">])</span>

    <span class="k">if</span> <span class="n">async_op</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">work</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span></div>

<span class="k">def</span> <span class="nf">all_gather_coalesced</span><span class="p">(</span><span class="n">output_tensor_lists</span><span class="p">,</span>
                         <span class="n">input_tensor_list</span><span class="p">,</span>
                         <span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">,</span>
                         <span class="n">async_op</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Gathers input tensors from the whole group in a list in a coalesced manner.</span>

<span class="sd">    Arguments:</span>
<span class="sd">        output_tensor_lists (list[list[Tensor]]): Output list. It should contain</span>
<span class="sd">            correctly-sized tensors to be used for output of the collective.</span>
<span class="sd">        input_tensor_list (list[Tensor]): Tensors to be broadcast from</span>
<span class="sd">            current process. At least one tensor has to be non empty.</span>
<span class="sd">        group (ProcessGroup, optional): The process group to work on</span>
<span class="sd">        async_op (bool, optional): Whether this op should be an async op.</span>

<span class="sd">    Returns:</span>
<span class="sd">        Async work handle, if async_op is set to True.</span>
<span class="sd">        None, if not async_op or if not part of the group</span>

<span class="sd">    Example:</span>
<span class="sd">        we have 2 process groups, 2 ranks.</span>
<span class="sd">        rank 0 passes:</span>
<span class="sd">            input_tensor_list = [[[1, 1], [1, 1]], [2], [3, 3]]</span>
<span class="sd">            output_tensor_lists =</span>
<span class="sd">               [[[[-1, -1], [-1, -1]], [-1], [-1, -1]],</span>
<span class="sd">                [[[-1, -1], [-1, -1]], [-1], [-1, -1]]]</span>
<span class="sd">        rank 1 passes:</span>
<span class="sd">            input_tensor_list = [[[3, 3], [3, 3]], [5], [1, 1]]</span>
<span class="sd">            output_tensor_lists =</span>
<span class="sd">               [[[[-1, -1], [-1, -1]], [-1], [-1, -1]],</span>
<span class="sd">                [[[-1, -1], [-1, -1]], [-1], [-1, -1]]]</span>
<span class="sd">        both rank 0 and 1 get:</span>
<span class="sd">            output_tensor_lists =</span>
<span class="sd">               [[[1, 1], [1, 1]], [2], [3, 3]],</span>
<span class="sd">                [[3, 3], [3, 3]], [5], [1, 1]]].</span>

<span class="sd">    WARNING: at this time individual shape checking is not implemented across nodes.</span>
<span class="sd">    For example, if the rank 0 node passes [torch.rand(4), torch.rand(2)] and the</span>
<span class="sd">    rank 1 node passes [torch.rand(2), torch.rand(2), torch.rand(2)], the</span>
<span class="sd">    all_gather_coalesced operation will proceed without complaint and return</span>
<span class="sd">    erroneous outputs. This lack of shape checking results in significant</span>
<span class="sd">    performance improvements but users of this function should take extra care</span>
<span class="sd">    to ensure that each node passes in tensors whose shapes match across nodes.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="c1"># We only check basic compatibility with C++ params here, C++ code will</span>
    <span class="c1"># do shape and type checking.</span>
    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span>
    <span class="n">_check_tensor_list</span><span class="p">(</span><span class="n">input_tensor_list</span><span class="p">,</span> <span class="s2">&quot;tensor_list&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">output_tensor_lists</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
        <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Invalid function argument: &quot;</span>
                     <span class="s2">&quot;output_tensor_lists should be a list&quot;</span><span class="p">)</span>
    <span class="k">for</span> <span class="n">output_tensor_list</span> <span class="ow">in</span> <span class="n">output_tensor_lists</span><span class="p">:</span>
        <span class="n">_check_tensor_list</span><span class="p">(</span><span class="n">output_tensor_list</span><span class="p">,</span> <span class="s2">&quot;output_tensor_lists&quot;</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_check_default_pg</span><span class="p">()</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">allgather_coalesced</span><span class="p">(</span>
            <span class="n">output_tensor_lists</span><span class="p">,</span> <span class="n">input_tensor_list</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">group</span><span class="o">.</span><span class="n">allgather_coalesced</span><span class="p">(</span><span class="n">output_tensor_lists</span><span class="p">,</span> <span class="n">input_tensor_list</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">async_op</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">work</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span>


<div class="viewcode-block" id="gather"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.gather">[docs]</a><span class="k">def</span> <span class="nf">gather</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span>
           <span class="n">gather_list</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
           <span class="n">dst</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
           <span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">,</span>
           <span class="n">async_op</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Gathers a list of tensors in a single process.</span>

<span class="sd">    Arguments:</span>
<span class="sd">        tensor (Tensor): Input tensor.</span>
<span class="sd">        gather_list (list[Tensor], optional): List of appropriately-sized</span>
<span class="sd">            tensors to use for gathered data (default is None, must be specified</span>
<span class="sd">            on the destination rank)</span>
<span class="sd">        dst (int, optional): Destination rank (default is 0)</span>
<span class="sd">        group (ProcessGroup, optional): The process group to work on</span>
<span class="sd">        async_op (bool, optional): Whether this op should be an async op</span>

<span class="sd">    Returns:</span>
<span class="sd">        Async work handle, if async_op is set to True.</span>
<span class="sd">        None, if not async_op or if not part of the group</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">_check_single_tensor</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="s2">&quot;tensor&quot;</span><span class="p">)</span>

    <span class="c1"># Parameter ``gather_list`` may be left unspecified on non-dst ranks.</span>
    <span class="k">if</span> <span class="n">gather_list</span><span class="p">:</span>
        <span class="n">_check_tensor_list</span><span class="p">(</span><span class="n">gather_list</span><span class="p">,</span> <span class="s2">&quot;gather_list&quot;</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">gather_list</span> <span class="o">=</span> <span class="p">[]</span>

    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span>

    <span class="n">my_rank</span> <span class="o">=</span> <span class="n">get_rank</span><span class="p">()</span>
    <span class="k">if</span> <span class="n">dst</span> <span class="o">==</span> <span class="n">my_rank</span><span class="p">:</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">gather_list</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Argument ``gather_list`` must be specified &quot;</span>
                             <span class="s2">&quot;on destination rank.&quot;</span><span class="p">)</span>
        <span class="n">input_tensors</span> <span class="o">=</span> <span class="p">[</span><span class="n">tensor</span><span class="p">]</span>
        <span class="n">output_tensors</span> <span class="o">=</span> <span class="p">[</span><span class="n">gather_list</span><span class="p">]</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">gather_list</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Argument ``gather_list`` must NOT be specified &quot;</span>
                             <span class="s2">&quot;on non-destination ranks.&quot;</span><span class="p">)</span>
        <span class="n">input_tensors</span> <span class="o">=</span> <span class="p">[</span><span class="n">tensor</span><span class="p">]</span>
        <span class="n">output_tensors</span> <span class="o">=</span> <span class="p">[]</span>

    <span class="n">opts</span> <span class="o">=</span> <span class="n">GatherOptions</span><span class="p">()</span>
    <span class="n">opts</span><span class="o">.</span><span class="n">rootRank</span> <span class="o">=</span> <span class="n">dst</span>

    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_check_default_pg</span><span class="p">()</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">gather</span><span class="p">(</span><span class="n">output_tensors</span><span class="p">,</span> <span class="n">input_tensors</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">group_dst_rank</span> <span class="o">=</span> <span class="n">_get_group_rank</span><span class="p">(</span><span class="n">group</span><span class="p">,</span> <span class="n">dst</span><span class="p">)</span>
        <span class="n">opts</span><span class="o">.</span><span class="n">rootRank</span> <span class="o">=</span> <span class="n">group_dst_rank</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">group</span><span class="o">.</span><span class="n">gather</span><span class="p">(</span><span class="n">output_tensors</span><span class="p">,</span> <span class="n">input_tensors</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">async_op</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">work</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span></div>


<div class="viewcode-block" id="scatter"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.scatter">[docs]</a><span class="k">def</span> <span class="nf">scatter</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span>
            <span class="n">scatter_list</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
            <span class="n">src</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
            <span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">,</span>
            <span class="n">async_op</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Scatters a list of tensors to all processes in a group.</span>

<span class="sd">    Each process will receive exactly one tensor and store its data in the</span>
<span class="sd">    ``tensor`` argument.</span>

<span class="sd">    Arguments:</span>
<span class="sd">        tensor (Tensor): Output tensor.</span>
<span class="sd">        scatter_list (list[Tensor]): List of tensors to scatter (default is</span>
<span class="sd">            None, must be specified on the source rank)</span>
<span class="sd">        src (int): Source rank (default is 0)</span>
<span class="sd">        group (ProcessGroup, optional): The process group to work on</span>
<span class="sd">        async_op (bool, optional): Whether this op should be an async op</span>

<span class="sd">    Returns:</span>
<span class="sd">        Async work handle, if async_op is set to True.</span>
<span class="sd">        None, if not async_op or if not part of the group</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">_check_single_tensor</span><span class="p">(</span><span class="n">tensor</span><span class="p">,</span> <span class="s2">&quot;tensor&quot;</span><span class="p">)</span>

    <span class="c1"># Parameter ``scatter_list`` may be left unspecified on non-src ranks.</span>
    <span class="k">if</span> <span class="n">scatter_list</span><span class="p">:</span>
        <span class="n">_check_tensor_list</span><span class="p">(</span><span class="n">scatter_list</span><span class="p">,</span> <span class="s2">&quot;scatter_list&quot;</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">scatter_list</span> <span class="o">=</span> <span class="p">[]</span>

    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span>

    <span class="n">my_rank</span> <span class="o">=</span> <span class="n">get_rank</span><span class="p">()</span>
    <span class="k">if</span> <span class="n">src</span> <span class="o">==</span> <span class="n">my_rank</span><span class="p">:</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="n">scatter_list</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Argument ``scatter_list`` must be specified &quot;</span>
                             <span class="s2">&quot;on source rank.&quot;</span><span class="p">)</span>
        <span class="n">input_tensors</span> <span class="o">=</span> <span class="p">[</span><span class="n">scatter_list</span><span class="p">]</span>
        <span class="n">output_tensors</span> <span class="o">=</span> <span class="p">[</span><span class="n">tensor</span><span class="p">]</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">scatter_list</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Argument ``scatter_list`` must NOT be specified &quot;</span>
                             <span class="s2">&quot;on non-source ranks.&quot;</span><span class="p">)</span>
        <span class="n">input_tensors</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="n">output_tensors</span> <span class="o">=</span> <span class="p">[</span><span class="n">tensor</span><span class="p">]</span>

    <span class="n">opts</span> <span class="o">=</span> <span class="n">ScatterOptions</span><span class="p">()</span>
    <span class="n">opts</span><span class="o">.</span><span class="n">rootRank</span> <span class="o">=</span> <span class="n">src</span>

    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_check_default_pg</span><span class="p">()</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span><span class="n">output_tensors</span><span class="p">,</span> <span class="n">input_tensors</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">group_src_rank</span> <span class="o">=</span> <span class="n">_get_group_rank</span><span class="p">(</span><span class="n">group</span><span class="p">,</span> <span class="n">src</span><span class="p">)</span>
        <span class="n">opts</span><span class="o">.</span><span class="n">rootRank</span> <span class="o">=</span> <span class="n">group_src_rank</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">group</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span><span class="n">output_tensors</span><span class="p">,</span> <span class="n">input_tensors</span><span class="p">,</span> <span class="n">opts</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">async_op</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">work</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span></div>


<span class="k">def</span> <span class="nf">reduce_scatter_multigpu</span><span class="p">(</span><span class="n">output_tensor_list</span><span class="p">,</span>
                            <span class="n">input_tensor_lists</span><span class="p">,</span>
                            <span class="n">op</span><span class="o">=</span><span class="n">ReduceOp</span><span class="o">.</span><span class="n">SUM</span><span class="p">,</span>
                            <span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">,</span>
                            <span class="n">async_op</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Reduce and scatter a list of tensors to the whole group.  Only nccl backend</span>
<span class="sd">    is currently supported.</span>

<span class="sd">    Each tensor in ``output_tensor_list`` should reside on a separate GPU, as</span>
<span class="sd">    should each list of tensors in ``input_tensor_lists``.</span>

<span class="sd">    Arguments:</span>
<span class="sd">        output_tensor_list (List[Tensor]): Output tensors (on different GPUs)</span>
<span class="sd">            to receive the result of the operation.</span>

<span class="sd">            Note that ``len(output_tensor_list)`` needs to be the same for all</span>
<span class="sd">            the distributed processes calling this function.</span>

<span class="sd">        input_tensor_lists (List[List[Tensor]]): Input lists.  It should</span>
<span class="sd">            contain correctly-sized tensors on each GPU to be used for input of</span>
<span class="sd">            the collective, e.g. ``input_tensor_lists[i]`` contains the</span>
<span class="sd">            reduce_scatter input that resides on the GPU of</span>
<span class="sd">            ``output_tensor_list[i]``.</span>

<span class="sd">            Note that each element of ``input_tensor_lists`` has the size of</span>
<span class="sd">            ``world_size * len(output_tensor_list)``, since the function</span>
<span class="sd">            scatters the result from every single GPU in the group.  To</span>
<span class="sd">            interpret each element of ``input_tensor_lists[i]``, note that</span>
<span class="sd">            ``output_tensor_list[j]`` of rank k receives the reduce-scattered</span>
<span class="sd">            result from ``input_tensor_lists[i][k * world_size + j]``</span>

<span class="sd">            Also note that ``len(input_tensor_lists)``, and the size of each</span>
<span class="sd">            element in ``input_tensor_lists`` (each element is a list,</span>
<span class="sd">            therefore ``len(input_tensor_lists[i])``) need to be the same for</span>
<span class="sd">            all the distributed processes calling this function.</span>

<span class="sd">        group (ProcessGroup, optional): The process group to work on.</span>
<span class="sd">        async_op (bool, optional): Whether this op should be an async op.</span>

<span class="sd">    Returns:</span>
<span class="sd">        Async work handle, if async_op is set to True.</span>
<span class="sd">        None, if not async_op or if not part of the group.</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span>

    <span class="n">opts</span> <span class="o">=</span> <span class="n">ReduceScatterOptions</span><span class="p">()</span>
    <span class="n">opts</span><span class="o">.</span><span class="n">reduceOp</span> <span class="o">=</span> <span class="n">op</span>

    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_check_default_pg</span><span class="p">()</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">reduce_scatter</span><span class="p">(</span>
            <span class="n">output_tensor_list</span><span class="p">,</span>
            <span class="n">input_tensor_lists</span><span class="p">,</span>
            <span class="n">opts</span>
        <span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">group</span><span class="o">.</span><span class="n">reduce_scatter</span><span class="p">(</span>
            <span class="n">output_tensor_list</span><span class="p">,</span>
            <span class="n">input_tensor_lists</span><span class="p">,</span>
            <span class="n">opts</span>
        <span class="p">)</span>

    <span class="k">if</span> <span class="n">async_op</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">work</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span>


<span class="k">def</span> <span class="nf">reduce_scatter</span><span class="p">(</span><span class="n">output</span><span class="p">,</span>
                   <span class="n">input_list</span><span class="p">,</span>
                   <span class="n">op</span><span class="o">=</span><span class="n">ReduceOp</span><span class="o">.</span><span class="n">SUM</span><span class="p">,</span>
                   <span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">,</span>
                   <span class="n">async_op</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Reduces, then scatters a list of tensors to all processes in a group.</span>

<span class="sd">    Arguments:</span>
<span class="sd">        output (Tensor): Output tensor.</span>
<span class="sd">        input_list (list[Tensor]): List of tensors to reduce and scatter.</span>
<span class="sd">        group (ProcessGroup, optional): The process group to work on.</span>
<span class="sd">        async_op (bool, optional): Whether this op should be an async op.</span>

<span class="sd">    Returns:</span>
<span class="sd">        Async work handle, if async_op is set to True.</span>
<span class="sd">        None, if not async_op or if not part of the group.</span>

<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">_check_single_tensor</span><span class="p">(</span><span class="n">output</span><span class="p">,</span> <span class="s2">&quot;output&quot;</span><span class="p">)</span>
    <span class="n">_check_tensor_list</span><span class="p">(</span><span class="n">input_list</span><span class="p">,</span> <span class="s2">&quot;input_list&quot;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span>

    <span class="n">opts</span> <span class="o">=</span> <span class="n">ReduceScatterOptions</span><span class="p">()</span>
    <span class="n">opts</span><span class="o">.</span><span class="n">reduceOp</span> <span class="o">=</span> <span class="n">op</span>

    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_check_default_pg</span><span class="p">()</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">reduce_scatter</span><span class="p">([</span><span class="n">output</span><span class="p">],</span> <span class="p">[</span><span class="n">input_list</span><span class="p">],</span> <span class="n">opts</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">group</span><span class="o">.</span><span class="n">reduce_scatter</span><span class="p">([</span><span class="n">output</span><span class="p">],</span> <span class="p">[</span><span class="n">input_list</span><span class="p">],</span> <span class="n">opts</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">async_op</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">work</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span>


<div class="viewcode-block" id="barrier"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.barrier">[docs]</a><span class="k">def</span> <span class="nf">barrier</span><span class="p">(</span><span class="n">group</span><span class="o">=</span><span class="n">group</span><span class="o">.</span><span class="n">WORLD</span><span class="p">,</span>
            <span class="n">async_op</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Synchronizes all processes.</span>

<span class="sd">    This collective blocks processes until the whole group enters this function,</span>
<span class="sd">    if async_op is False, or if async work handle is called on wait().</span>

<span class="sd">    Arguments:</span>
<span class="sd">        group (ProcessGroup, optional): The process group to work on</span>
<span class="sd">        async_op (bool, optional): Whether this op should be an async op</span>

<span class="sd">    Returns:</span>
<span class="sd">        Async work handle, if async_op is set to True.</span>
<span class="sd">        None, if not async_op or if not part of the group</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">if</span> <span class="n">_rank_not_in_group</span><span class="p">(</span><span class="n">group</span><span class="p">):</span>
        <span class="k">return</span>

    <span class="k">if</span> <span class="n">group</span> <span class="o">==</span> <span class="n">GroupMember</span><span class="o">.</span><span class="n">WORLD</span><span class="p">:</span>
        <span class="n">_check_default_pg</span><span class="p">()</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">barrier</span><span class="p">()</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span> <span class="o">=</span> <span class="n">group</span><span class="o">.</span><span class="n">barrier</span><span class="p">()</span>

    <span class="k">if</span> <span class="n">async_op</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">work</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">work</span><span class="o">.</span><span class="n">wait</span><span class="p">()</span></div>


<div class="viewcode-block" id="new_group"><a class="viewcode-back" href="../../../distributed.html#torch.distributed.new_group">[docs]</a><span class="k">def</span> <span class="nf">new_group</span><span class="p">(</span><span class="n">ranks</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">timeout</span><span class="o">=</span><span class="n">default_pg_timeout</span><span class="p">,</span> <span class="n">backend</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Creates a new distributed group.</span>

<span class="sd">    This function requires that all processes in the main group (i.e. all</span>
<span class="sd">    processes that are part of the distributed job) enter this function, even</span>
<span class="sd">    if they are not going to be members of the group. Additionally, groups</span>
<span class="sd">    should be created in the same order in all processes.</span>

<span class="sd">    Arguments:</span>
<span class="sd">        ranks (list[int]): List of ranks of group members.</span>
<span class="sd">        timeout (timedelta, optional): Timeout for operations executed against</span>
<span class="sd">            the process group. Default value equals 30 minutes.</span>
<span class="sd">            This is only applicable for the ``gloo`` backend.</span>
<span class="sd">        backend (str or Backend, optional): The backend to use. Depending on</span>
<span class="sd">            build-time configurations, valid values are ``gloo`` and ``nccl``.</span>
<span class="sd">            By default uses the same backend as the global group. This field</span>
<span class="sd">            should be given as a lowercase string (e.g., ``&quot;gloo&quot;``), which can</span>
<span class="sd">            also be accessed via :class:`Backend` attributes (e.g.,</span>
<span class="sd">            ``Backend.GLOO``).</span>

<span class="sd">    Returns:</span>
<span class="sd">        A handle of distributed group that can be given to collective calls.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="n">_check_default_pg</span><span class="p">()</span>

    <span class="k">global</span> <span class="n">_pg_group_ranks</span>

    <span class="n">default_backend</span><span class="p">,</span> <span class="n">default_store</span> <span class="o">=</span> <span class="n">_pg_map</span><span class="p">[</span><span class="n">_default_pg</span><span class="p">]</span>
    <span class="n">global_rank</span> <span class="o">=</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">rank</span><span class="p">()</span>
    <span class="n">global_world_size</span> <span class="o">=</span> <span class="n">_default_pg</span><span class="o">.</span><span class="n">size</span><span class="p">()</span>

    <span class="c1"># Default to the same backend as the global process group</span>
    <span class="c1"># if the backend is not specified.</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">backend</span><span class="p">:</span>
        <span class="n">backend</span> <span class="o">=</span> <span class="n">default_backend</span>

    <span class="c1"># checks the input ranks</span>
    <span class="k">if</span> <span class="n">ranks</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">ranks</span> <span class="o">=</span> <span class="nb">sorted</span><span class="p">(</span><span class="n">ranks</span><span class="p">)</span>
        <span class="n">group_world_size</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">ranks</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">group_world_size</span> <span class="o">&gt;</span> <span class="n">global_world_size</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;the new group&#39;s world size should be less or &quot;</span>
                               <span class="s2">&quot;equal to the world size set by &quot;</span>
                               <span class="s2">&quot;init_process_group&quot;</span><span class="p">)</span>
        <span class="c1"># check ranks&#39; sanity</span>
        <span class="k">for</span> <span class="n">rank</span> <span class="ow">in</span> <span class="n">ranks</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">rank</span> <span class="o">&lt;</span> <span class="mi">0</span> <span class="ow">or</span> <span class="n">rank</span> <span class="o">&gt;=</span> <span class="n">global_world_size</span><span class="p">:</span>
                <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;The new group&#39;s rank should be within the &quot;</span>
                                   <span class="s2">&quot;the world_size set by init_process_group&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">global_rank</span> <span class="ow">in</span> <span class="n">ranks</span><span class="p">:</span>
            <span class="n">group_rank</span> <span class="o">=</span> <span class="n">ranks</span><span class="o">.</span><span class="n">index</span><span class="p">(</span><span class="n">global_rank</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">group_rank</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">ranks</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="n">global_world_size</span><span class="p">))</span>
        <span class="n">group_world_size</span> <span class="o">=</span> <span class="n">global_world_size</span>
        <span class="n">group_rank</span> <span class="o">=</span> <span class="n">global_rank</span>

    <span class="n">backend</span> <span class="o">=</span> <span class="n">Backend</span><span class="p">(</span><span class="n">backend</span><span class="p">)</span>
    <span class="n">pg</span> <span class="o">=</span> <span class="n">_new_process_group_helper</span><span class="p">(</span><span class="n">group_world_size</span><span class="p">,</span>
                                   <span class="n">group_rank</span><span class="p">,</span>
                                   <span class="n">ranks</span><span class="p">,</span>
                                   <span class="n">backend</span><span class="p">,</span>
                                   <span class="n">default_store</span><span class="p">,</span>
                                   <span class="n">timeout</span><span class="o">=</span><span class="n">timeout</span><span class="p">)</span>

    <span class="c1"># Create the global rank to group rank mapping</span>
    <span class="n">_pg_group_ranks</span><span class="p">[</span><span class="n">pg</span><span class="p">]</span> <span class="o">=</span> <span class="p">{</span>
        <span class="n">global_rank</span><span class="p">:</span> <span class="n">group_rank</span>
        <span class="k">for</span> <span class="n">group_rank</span><span class="p">,</span> <span class="n">global_rank</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">ranks</span><span class="p">)</span>
    <span class="p">}</span>

    <span class="k">return</span> <span class="n">pg</span></div>
</pre></div>

             </article>
             
            </div>
            <footer>
  

  

    <hr>

  

  <div role="contentinfo">
    <p>
        &copy; Copyright 2019, Torch Contributors.

    </p>
  </div>
    
      <div>
        Built with <a href="http://sphinx-doc.org/">Sphinx</a> using a <a href="https://github.com/rtfd/sphinx_rtd_theme">theme</a> provided by <a href="https://readthedocs.org">Read the Docs</a>.
      </div>
     

</footer>

          </div>
        </div>

        <div class="pytorch-content-right" id="pytorch-content-right">
          <div class="pytorch-right-menu" id="pytorch-right-menu">
            <div class="pytorch-side-scroll" id="pytorch-side-scroll-right">
              
            </div>
          </div>
        </div>
      </section>
    </div>

  


  

     
       <script type="text/javascript" id="documentation_options" data-url_root="../../../" src="../../../_static/documentation_options.js"></script>
         <script src="../../../_static/jquery.js"></script>
         <script src="../../../_static/underscore.js"></script>
         <script src="../../../_static/doctools.js"></script>
         <script src="../../../_static/language_data.js"></script>
     

  

  <script type="text/javascript" src="../../../_static/js/vendor/popper.min.js"></script>
  <script type="text/javascript" src="../../../_static/js/vendor/bootstrap.min.js"></script>
  <script src="https://cdnjs.cloudflare.com/ajax/libs/list.js/1.5.0/list.min.js"></script>
  <script type="text/javascript" src="../../../_static/js/theme.js"></script>

  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>
 
<script>
  (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){
  (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o),
  m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m)
  })(window,document,'script','https://www.google-analytics.com/analytics.js','ga');

  ga('create', 'UA-90545585-1', 'auto');
  ga('send', 'pageview');

</script>

<script async src="https://www.googletagmanager.com/gtag/js?id=UA-117752657-2"></script>

<script>
  window.dataLayer = window.dataLayer || [];

  function gtag(){dataLayer.push(arguments);}

  gtag('js', new Date());
  gtag('config', 'UA-117752657-2');
</script>

<img height="1" width="1" style="border-style:none;" alt="" src="https://www.googleadservices.com/pagead/conversion/795629140/?label=txkmCPmdtosBENSssfsC&amp;guid=ON&amp;script=0"/>


  <!-- Begin Footer -->

  <div class="container-fluid docs-tutorials-resources" id="docs-tutorials-resources">
    <div class="container">
      <div class="row">
        <div class="col-md-4 text-center">
          <h2>Docs</h2>
          <p>Access comprehensive developer documentation for PyTorch</p>
          <a class="with-right-arrow" href="https://pytorch.org/docs/stable/index.html">View Docs</a>
        </div>

        <div class="col-md-4 text-center">
          <h2>Tutorials</h2>
          <p>Get in-depth tutorials for beginners and advanced developers</p>
          <a class="with-right-arrow" href="https://pytorch.org/tutorials">View Tutorials</a>
        </div>

        <div class="col-md-4 text-center">
          <h2>Resources</h2>
          <p>Find development resources and get your questions answered</p>
          <a class="with-right-arrow" href="https://pytorch.org/resources">View Resources</a>
        </div>
      </div>
    </div>
  </div>

  <footer class="site-footer">
    <div class="container footer-container">
      <div class="footer-logo-wrapper">
        <a href="https://pytorch.org/" class="footer-logo"></a>
      </div>

      <div class="footer-links-wrapper">
        <div class="footer-links-col">
          <ul>
            <li class="list-title"><a href="https://pytorch.org/">PyTorch</a></li>
            <li><a href="https://pytorch.org/get-started">Get Started</a></li>
            <li><a href="https://pytorch.org/features">Features</a></li>
            <li><a href="https://pytorch.org/ecosystem">Ecosystem</a></li>
            <li><a href="https://pytorch.org/blog/">Blog</a></li>
            <li><a href="https://github.com/pytorch/pytorch/blob/master/CONTRIBUTING.md">Contributing</a></li>
          </ul>
        </div>

        <div class="footer-links-col">
          <ul>
            <li class="list-title"><a href="https://pytorch.org/resources">Resources</a></li>
            <li><a href="https://pytorch.org/tutorials">Tutorials</a></li>
            <li><a href="https://pytorch.org/docs/stable/index.html">Docs</a></li>
            <li><a href="https://discuss.pytorch.org" target="_blank">Discuss</a></li>
            <li><a href="https://github.com/pytorch/pytorch/issues" target="_blank">Github Issues</a></li>
            <li><a href="https://pytorch.org/assets/brand-guidelines/PyTorch-Brand-Guidelines.pdf" target="_blank">Brand Guidelines</a></li>
          </ul>
        </div>

        <div class="footer-links-col follow-us-col">
          <ul>
            <li class="list-title">Stay Connected</li>
            <li>
              <div id="mc_embed_signup">
                <form
                  action="https://twitter.us14.list-manage.com/subscribe/post?u=75419c71fe0a935e53dfa4a3f&id=91d0dccd39"
                  method="post"
                  id="mc-embedded-subscribe-form"
                  name="mc-embedded-subscribe-form"
                  class="email-subscribe-form validate"
                  target="_blank"
                  novalidate>
                  <div id="mc_embed_signup_scroll" class="email-subscribe-form-fields-wrapper">
                    <div class="mc-field-group">
                      <label for="mce-EMAIL" style="display:none;">Email Address</label>
                      <input type="email" value="" name="EMAIL" class="required email" id="mce-EMAIL" placeholder="Email Address">
                    </div>

                    <div id="mce-responses" class="clear">
                      <div class="response" id="mce-error-response" style="display:none"></div>
                      <div class="response" id="mce-success-response" style="display:none"></div>
                    </div>    <!-- real people should not fill this in and expect good things - do not remove this or risk form bot signups-->

                    <div style="position: absolute; left: -5000px;" aria-hidden="true"><input type="text" name="b_75419c71fe0a935e53dfa4a3f_91d0dccd39" tabindex="-1" value=""></div>

                    <div class="clear">
                      <input type="submit" value="" name="subscribe" id="mc-embedded-subscribe" class="button email-subscribe-button">
                    </div>
                  </div>
                </form>
              </div>

            </li>
          </ul>

          <div class="footer-social-icons">
            <a href="https://www.facebook.com/pytorch" target="_blank" class="facebook"></a>
            <a href="https://twitter.com/pytorch" target="_blank" class="twitter"></a>
            <a href="https://www.youtube.com/pytorch" target="_blank" class="youtube"></a>
          </div>
        </div>
      </div>
    </div>
  </footer>

  <div class="cookie-banner-wrapper">
  <div class="container">
    <p class="gdpr-notice">To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: <a href="https://www.facebook.com/policies/cookies/">Cookies Policy</a>.</p>
    <img class="close-button" src="../../../_static/images/pytorch-x.svg">
  </div>
</div>

  <!-- End Footer -->

  <!-- Begin Mobile Menu -->

  <div class="mobile-main-menu">
    <div class="container-fluid">
      <div class="container">
        <div class="mobile-main-menu-header-container">
          <a class="header-logo" href="https://pytorch.org/" aria-label="PyTorch"></a>
          <a class="main-menu-close-button" href="#" data-behavior="close-mobile-menu"></a>
        </div>
      </div>
    </div>

    <div class="mobile-main-menu-links-container">
      <div class="main-menu">
        <ul>
          <li>
            <a href="https://pytorch.org/get-started">Get Started</a>
          </li>

          <li>
            <a href="https://pytorch.org/features">Features</a>
          </li>

          <li>
            <a href="https://pytorch.org/ecosystem">Ecosystem</a>
          </li>

          <li>
            <a href="https://pytorch.org/mobile">Mobile</a>
          </li>

          <li>
            <a href="https://pytorch.org/hub">PyTorch Hub</a>
          </li>

          <li>
            <a href="https://pytorch.org/blog/">Blog</a>
          </li>

          <li>
            <a href="https://pytorch.org/tutorials">Tutorials</a>
          </li>

          <li class="active">
            <a href="https://pytorch.org/docs/stable/index.html">Docs</a>
          </li>

          <li>
            <a href="https://pytorch.org/resources">Resources</a>
          </li>

          <li>
            <a href="https://github.com/pytorch/pytorch">Github</a>
          </li>
        </ul>
      </div>
    </div>
  </div>

  <!-- End Mobile Menu -->

  <script type="text/javascript" src="../../../_static/js/vendor/anchor.min.js"></script>

  <script type="text/javascript">
    $(document).ready(function() {
      mobileMenu.bind();
      mobileTOC.bind();
      pytorchAnchors.bind();
      sideMenus.bind();
      scrollToAnchor.bind();
      highlightNavigation.bind();
      mainMenuDropdown.bind();
      filterTags.bind();

      // Remove any empty p tags that Sphinx adds
      $("[data-tags='null']").remove();

      // Add class to links that have code blocks, since we cannot create links in code blocks
      $("article.pytorch-article a span.pre").each(function(e) {
        $(this).closest("a").addClass("has-code");
      });
    })
  </script>
</body>
</html>